How to feed both the memory state and the carry state to a LSTMCell in Tensorflow 1.x?

Viewed 41

I'm trying to reproduce an experiment using Tensorflow 1.14. It built a Graph Neural Network using GRUCell as below:

self.edge_update = tf.keras.layers.GRUCell(hparams.link_state_dim)
self.path_update = tf.keras.layers.GRUCell(hparams.path_state_dim)

And it called the function like this:

...
outputs, path_state = tf.compat.v1.nn.dynamic_rnn(self.path_update,
                                        link_inputs,
                                        sequence_length=lens,
                                        initial_state = path_state,
                                        dtype=tf.float32)
...

link_state,_ = self.edge_update(m, [link_state])

Note that the path_state and link_state is self-definded.

The code is running well while using the initial GRUCell. And now I want to replace it with LSTMCell.

self.edge_update = tf.keras.layers.LSTMCell(hparams.link_state_dim)
self.path_update = tf.keras.layers.LSTMCell(hparams.path_state_dim)

Considering that LSTM has two states named memory state and carry state respectively, I tried to feed the self-defined initial state to both of them. So I tried as below:

...
outputs, path_state = tf.compat.v1.nn.dynamic_rnn(self.path_update,
                                        link_inputs,
                                        sequence_length=lens,
                                        initial_state = [path_state,path_state],
                                        dtype=tf.float32)
...
link_state,_ = self.edge_update(m, [link_state,link_state])

But it turned out that this is illegal. The error was reported like this:

AttributeError: 'list' object has no attribute 'shape'

And I have also checked the official documentation.It gives the calling method of both GRUCell and LSTMCell like below:

GRUCell: states: A 2D tensor with shape of [batch, units], which is the state from the previous time step. For timestep 0, the initial state provided by user will be feed to cell.

LSTMCell: states: List of 2 tensors that corresponding to the cell's units. Both of them have shape [batch, units], the first tensor is the memory state from previous time step, the second tensor is the carry state from previous time step. For timestep 0, the initial state provided by user will be feed to cell.

It seems that I had followed the official documents, but it's invalid. So how can I feed both memory state and carry state to a LSTMCell?

0 Answers
Related