首页 > 解决方案 > 初始化 CUDNN LSTM 的状态

问题描述

我认为我们可以使用以下代码段来创建 LSTM 堆栈并将其状态初始化为零。

 lstm_cell = tf.contrib.rnn.BasicLSTMCell(
            hidden_size, forget_bias=0.0, state_is_tuple=True)
 cell = tf.contrib.rnn.MultiRNNCell([lstm_cell] * num_layers, state_is_tuple=True)
 cell.zero_state(batch_size, tf_float32)

我想使用 CUDNN 而不是使用 BasicLSTMCell

cudnn_cell = tf.contrib.cudnn_rnn.CudnnLSTM(
          num_layers, hidden_size, dropout=config.keep_prob)

在这种情况下,我该如何做与cell.zero_state(batch_size, tf_float32)cudnn_cell 相同的事情?

标签: tensorflowstatecudnn

解决方案


定义见:tensorflow cudnn_rnn's code

关于 initial_states:

with tf.Graph().as_default():
    lstm = CudnnLSTM(num_layers, num_units, direction, ...)
    outputs, output_states = lstm(inputs, initial_states, training=True)

因此,除了嵌入输入之外,您只需要添加初始状态即可。在编码器-解码器系统中,它看起来像:

encoder_cell = tf.contrib.cudnn_rnn.CudnnLSTM(num_layers, hidden_size)
encoder_output, encoder_state = encoder_cell(encoder_embedding_input)
decoder_cell = tf.contrib.cudnn_rnn.CudnnLSTM(num_layers, hidden_size)
decoder_output, decoder_state = encoder_cell(decoder_embedding_input,
                                             initial_states=encoder_state)

在这里,encoder_state是一个tupleas (final_c_state, final_h_state)。两种状态的形状都是(1, batch, hidden_size)

如果您的编码器是双向 RNN,那将有点棘手,因为输出状态现在变为 (2, batch, hidden_size).

因此,我使用迂回的方式来解决它。

encoder_cell = tf.contrib.cudnn_rnn.CudnnLSTM(num_layers, hidden_size, direction="bidirectional")
encoder_output, (encoder_c_state, encoder_h_state) = encoder_cell(encoder_embedding_input)
fw_c, bw_c = tf.split(encoder_c_state, [1, 1], axis=0)
fw_h, bw_h = tf.split(encoder_h_state, [1, 1], axis=0)
reshape_encoder_c_state = tf.concat((fw_c, bw_c), axis=2)
reshape_encoder_h_state = tf.concat((fw_h, bw_h), axis=2)
encoder_state = (reshape_encoder_c_state, reshape_encoder_h_state)

虽然没试过多层RNN,但我觉得也可以用类似的方法解决。


推荐阅读