tensorflow 双向RNN/LSTM代码实现

ゞ浴缸里的玫瑰 2022-11-16 05:45 206阅读 0赞

def my_lstm_layer(input_representation, 
      lstm_dim=100, 
      input_lengths=None, 
      scope_name="my_rnn", 
      reuse=False,
      is_training=True,
      dropout_rate=0.2, 
      use_cudnn=False):
        ''' input_representation: [batch_size, seq_len, feature_dim] input_lengths: [batch_size] （可以为空） '''
        if is_training:
            input_representation = tf.nn.dropout(input_representation, (1 - dropout_rate))
        with tf.variable_scope(scope_name, reuse=reuse):
            if use_cudnn:
                inputs = tf.transpose(input_representation, [1, 0, 2])
                lstm = tf.contrib.cudnn_rnn.CudnnLSTM(1, lstm_dim, direction="bidirectional",
                                                      name="{}_cudnn_bi_lstm".format(scope_name),
                                                      dropout=dropout_rate if is_training else 0)
                outputs, _ = lstm(inputs)
                outputs = tf.transpose(outputs, [1, 0, 2])
                forward_representation = outputs[:, :, 0:lstm_dim]
                backward_representation = outputs[:, :, lstm_dim:2 * lstm_dim]
            else:
                context_lstm_cell_fw = tf.nn.rnn_cell.BasicLSTMCell(lstm_dim)
                context_lstm_cell_bw = tf.nn.rnn_cell.BasicLSTMCell(lstm_dim)
                if is_training:
                    context_lstm_cell_fw = tf.nn.rnn_cell.DropoutWrapper(context_lstm_cell_fw,
                                                                         output_keep_prob=(1 - dropout_rate))
                    context_lstm_cell_bw = tf.nn.rnn_cell.DropoutWrapper(context_lstm_cell_bw,
                                                                         output_keep_prob=(1 - dropout_rate))
                context_lstm_cell_fw = tf.nn.rnn_cell.MultiRNNCell([context_lstm_cell_fw])
                context_lstm_cell_bw = tf.nn.rnn_cell.MultiRNNCell([context_lstm_cell_bw])
    
                (forward_representation, backward_representation), _ = tf.nn.bidirectional_dynamic_rnn(
                    context_lstm_cell_fw, context_lstm_cell_bw, input_representation, dtype=tf.float32,
                    sequence_length=input_lengths) 
                outputs = tf.concat(axis=2, values=[forward_representation, backward_representation])
        return (forward_representation, backward_representation, outputs)