in TensorFlow 2.0, what is the different between LSTM and LSTMCell objects?

Question

I am trying to implement an RNN in TensorFlow 2.0 (beta1). Looking at the layer functions (inherited from Keras) I found:

tf.keras.layers.LSTM

and

tf.keras.layers.LSTMCell

What is the different between the two? If you take a look at their arguments, they seem the same.

Answer 1

You are right - the difference is minimal. The base LSTMCell class implements the main functionality required, such as the build method, whereas the LSTM class only container an entry point: the call method, as well as a bunch of getters to retrieve attribute values. LSTMCell is the base class, which is used as a cell that is used inside the LSTM class.

All links point to relevant parts of tensorflow.keras source code.

My advice would be to use the standard LSTM class in your model as a normal layer. If you have a GPU at your disposal, you might want to use the version of the layer that is optimised with CUDA for execution on a GPU. As per the documentation:

Note that this cell is not optimized for performance on GPU. Please use tf.keras.layers.CuDNNLSTM for better performance on GPU.

There is also a GRU layer as well as a CuDNNGRU layer.

---

If you want to tweak how things work under the hood, you might create a class and inherit from the LSTMCell, or even the base class:

from tensorflow.python.keras.engine.base_layer import Layer

class MyLSTM(Layer):
    pass

But you would have to implement many things for yourself.