Howard Wang
Howard Wang

Reputation: 472

LSTM InvalidArgumentError Tensorflow 2.0/Keras when converting to estimator

I'm trying to build an LSTM network that takes in a sequence of words and convert it into an embedding vector. I've already converted each sequence of words into vocabulary vectors.

The batch size I'm using is 32 and the size of each vocabulary vector is 50. This is the Keras Functional API code I have so far to create the model and convert it to an estimator

input_layer = keras.layers.Input(shape=(50,), name='search')
embedding_layer = keras.layers.Embedding(input_dim=32, output_dim=256, input_length=50)(input_layer)
lstm_layer = keras.layers.LSTM(units=256)(embedding_layer)
model = keras.models.Model(inputs=input_layer, outputs=lstm_layer)
model.compile(loss='mean_squared_error', optimizer='adam')
estimator = keras.estimator.model_to_estimator(keras_model=model)

But this code gives the error

tensorflow.python.framework.errors_impl.InvalidArgumentError: Node 'Adam/gradients/lstm/StatefulPartitionedCall_grad/StatefulPartitionedCall': Connecting to invalid output 5 of source node lstm/StatefulPartitionedCall which has 5 outputs

When I run model.summary(), this is the output

Layer (type)                 Output Shape              Param #   
=================================================================
search (InputLayer)          [(None, 50)]              0         
_________________________________________________________________
embedding (Embedding)        (None, 50, 256)           8192      
_________________________________________________________________
lstm (LSTM)                  (None, 256)               525312    
=================================================================
Total params: 533,504
Trainable params: 533,504
Non-trainable params: 0
_________________________________________________________________

which I think is what I expect. I tried replacing the LSTM layer with a Dense and Flatten layer of the same shape and the code works fine

Upvotes: 1

Views: 476

Answers (1)

Howard Wang
Howard Wang

Reputation: 472

Going to answer this myself...it seems like there's problems with tf.keras.layers.LSTM as of 7/24 as seen here. I changed the model to the following

input_layer = keras.layers.Input(shape=(50,), name='search')
embedding_layer = keras.layers.Embedding(input_dim=32, output_dim=256, 
input_length=50)(input_layer)
lstm_layer = keras.layers.RNN(cell=keras.layers.LSTMCell(units=256))(embedding_layer)
model = keras.models.Model(inputs=input_layer, outputs=lstm_layer)
model.compile(loss='mean_squared_error', optimizer='adam')
estimator = keras.estimator.model_to_estimator(keras_model=model)

Upvotes: 1

Related Questions