Tfjs: Cannot load saved Keras model with recurrent layer

[psaikko]

TensorFlow.js version

@tensorflow/[email protected]/dist/tf.min.js

python packages:
tensorflow 2.2.0
tensorflowjs 1.7.4.post1

Browser version

Chrome: Version 81.0.4044.113 (Official Build) (64-bit)
Firefox: 76.0.1 (64-bit)

Describe the problem or feature request

When trying to load a saved Keras model with a recurrent layer in tensorflow.js, encountering
Error: Provided weight data has no target variable: simple_rnn/simple_rnn_cell/kernel or similar with LSTM/GRU layer.
Model is loaded correctly when recurrent layer is replaced with a dense layer

Code to reproduce the bug / link to feature request

Python:

import tensorflow as tf
import tensorflowjs as tfjs
model = tf.keras.Sequential()
model.add(tf.keras.layers.SimpleRNN(2, input_dim=2))
tfjs.converters.save_keras_model(model, 'gru_model')

html/js:

<script src="https://cdn.jsdelivr.net/npm/@tensorflow/[email protected]/dist/tf.min.js"></script>
<script>
tf.loadLayersModel('./gru_model/model.json').then(model => {
    model.summary();
})
</script>

[jiayihu]

I have a very similar error with a more complex RNN. The error, when I try to load the model in JS is:

Provided weight data has no target variable: bidirectional/forward_gru/gru_cell_1/kernel

The tensorflow keras model:

# GRU Encoder
encoder_in_layer = keras.layers.Input(shape=(max_length_in,))
encoder_embedding = keras.layers.Embedding(input_dim=vocab_size, output_dim=embedding_dim)
encoder_bi_lstm = keras.layers.Bidirectional(keras.layers.GRU(units=latent_dim, return_sequences=True, return_state=True))

# Discard the encoder output and use hidden states (h) and memory cells states (c)
# for forward (f) and backward (b) layer
encoder_out, fstate_h, bstate_h = encoder_bi_lstm(encoder_embedding(encoder_in_layer))
state_h = keras.layers.Concatenate()([fstate_h, bstate_h])

# GRUDecoder
decoder_in_layer = keras.layers.Input(shape=(None,))
decoder_embedding = keras.layers.Embedding(input_dim=vocab_size, output_dim=embedding_dim)
decoder_gru = keras.layers.GRU(units=latent_dim * 2, return_sequences=True, return_state=True)
# Discard internal states in training, keep only the output sequence
decoder_gru_out, _ = decoder_gru(decoder_embedding(decoder_in_layer), initial_state=state_h)
decoder_dense_1 = keras.layers.Dense(latent_dim, activation="relu")
decoder_dense = keras.layers.Dense(vocab_size, activation="softmax")
decoder_out_layer = decoder_dense(keras.layers.Dropout(rate=0.2)(decoder_dense_1(keras.layers.Dropout(rate=0.2)(decoder_gru_out))))

# Define the model that uses the Encoder and the Decoder
model = keras.models.Model([encoder_in_layer, decoder_in_layer], decoder_out_layer)

model.compile(optimizer='adam', loss="sparse_categorical_crossentropy", metrics=['sparse_categorical_accuracy'])
model.summary()

# Later...
encoder_model = keras.models.Model(encoder_in_layer, state_h)

And it's saved as h5 because I had issue converting the TF SaveModel with tensorflowjs-converter and I was able to do the conversion without errors only with the h5 model:

encoder_model.save('./encoder-model.h5')

The model is then loaded by just using tf.loadLayersModel('public/enc-model/model.json') in JavaScript. I tried passing { strict: false } to suppress the error but the resulting model has slightly less parameters the the original one and the predictions are completely wrong.

I have also tried to load the saved model using Python to check if the issue was on the saved h5 model itself and it works fine in Colab.

import tensorflow as tf
from tensorflow import keras

enc_model = keras.models.load_model('./encoder-model.h5', compile=False)

All the model and JS code can be found at https://github.com/jiayihu/gmail-smart-compose

I also checked the used versions. Colab is using Tensorflow 2.2, whereas I'm using Tensorflow.js 2.0 within the browser.

[psaikko]

I get the same Error: Provided weight data has no target variable: simple_rnn/simple_rnn_cell/kernel when saving as hdf5 and converting.

That is, model.save("./test-model.h5") in the python script followed by tensorflowjs_converter --input_format=keras test-model.h5 gru_model in bash, but the problem is the same.

[pyu10055]

I found the root cause of this, will submit a fix soon. @psaikko @jiayihu

[pyu10055]

There is still an issue with GRU layers shape mismatch for bias weight. Will investigate and fix that in a follow up PR.

[jiayihu]

@pyu10055 does it mean that it should work if I use LSTM units instead of GRU?

[psaikko]

I can confirm the diagnosis of @pyu10055 in https://github.com/tensorflow/tfjs/pull/3370.

Downgrading tensorflow from 2.2.0 to 2.1.0 resolved the problem on my end. Thanks!