Cannot load saved Keras model with recurrent layer

TensorFlow.js version

@tensorflow/tfjs@1.7.4/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/tfjs@1.7.4/dist/tf.min.js"></script>
<script>
tf.loadLayersModel('./gru_model/model.json').then(model => {
    model.summary();
})
</script>

[pyu10055]

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

[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.