如何使用Cudnn LSTM代替普通LSTM

Question

我一直在做一个命名实体识别任务。我在两者之间有以下代码段

with tf.variable_scope("Bi-lstm"):
        cell_fw = tf.contrib.rnn.LSTMCell(hidden_size)
        cell_bw = tf.contrib.rnn.LSTMCell(hidden_size)
        (output_fw, output_bw), _ = tf.nn.bidirectional_dynamic_rnn(cell_fw, 
            cell_bw, self.word_embeddings, sequence_length=self.sequence_lengths, 
            dtype=tf.float32)

现在我想使用Cudnn LSTM而不是普通的LSTM。在Cudnn LSTM的文档中，有一个额外的参数"num_layers“。我不明白这是什么，也不知道我该如何继续下去。

Answer 1

如果您使用的是TF，则不确定是哪个版本，但在TF 2.0文档中，您需要使用：

tf.keras.layers.LSTM

而不是：

tf.keras.layers.LSTMCell

https://www.tensorflow.org/guide/keras/rnn#using_cudnn_kernels_when_available