如何在 keras 中包装一个 tensorflow RNNCell？

Question

I would like to implement an custom LSTM cell as in a keras layer. Actually this implementation exists in tensorflow, so I was wondering if it is possible to just wrap it as an keras layer and call it in the model.

I found official documentation too simplified to see how to build a custom RNN layer. There are similar questions here and here , but they seem unresolved.

Thanks in advance for your help !

Answer 1

From my understanding, you should just be able to initialize the cell in the init() of the class layer, and then inside the call method reference it with your input.

Ex:

class MySimpleLayer(Layer):
  def __init__(self, lstm_size):
    super(MySimpleLayer, self).__init__()
    self.lstm = tf.contrib.rnn.BasicLSTMCell(lstm_size)

  def call(self, batch, state):
    return self.lstm(batch, state)

layer = MySimpleLayer(lstm_size)
logits = layer(batch, state)

This implementation is as basic as it gets though, so you might need to look into the build() and compute_output_shape() methods for more complex use cases.

Answer 2

Now the documentation of tensorflow may have improved since the question is posted.

You may want to check this guide or this SO answer for reference.