Tensorflow saving subclass model which has multiple arguments to call() method

Question

I am following the tensorflow neural machine translation tutorial: https://www.tensorflow.org/tutorials/text/nmt_with_attention

I am trying to save the Encoder and Decoder models which are subclasses of the tf.keras.Model and work properly during training and inference, however I want to save the models. When I try to do so I get the following error:

TypeError: call() missing 1 required positional argument: 'initial_state'

Here is the code:

class Encoder(tf.keras.Model):
    def __init__(self, vocab_size, embedding_matrix, n_units, batch_size):
        super(Encoder, self).__init__()
        self.n_units = n_units
        self.batch_size = batch_size

        self.embedding = Embedding(vocab_size, embedding_matrix.shape[1], weights=[embedding_matrix], trainable=True, mask_zero=True)
        self.lstm = LSTM(n_units, return_sequences=True, return_state=True, recurrent_initializer="glorot_uniform")

    def call(self, input_utterence, initial_state):
        input_embed = self.embedding(input_utterence)
        encoder_states, h1, c1 = self.lstm(input_embed, initial_state=initial_state)
        return encoder_states, h1, c1

    def create_initial_state(self):
        return tf.zeros((self.batch_size, self.n_units))

encoder = Encoder(vocab_size, embedding_matrix, LSTM_DIM, BATCH_SIZE)
# do some training...
tf.saved_model.save(decoder, "encoder_model")

I also tried to make the call method take one input list argument only and unpack the variables I need within the method but then I get the following error when trying to save:

File "C:\Users\Fady\Documents\Machine Learning\chatbot\models\seq2seq_model.py", line 32, in call
input_utterence, initial_state = inputs
ValueError: too many values to unpack (expected 2)

user_9579075 · Answer

You can export the model successfully if you package your inputs in a list. You also need to specify the input signatures to export your model, here your code with slight modifications which works

import tensorflow as tf
from tensorflow.keras.layers import Embedding, LSTM
import numpy as np
print('TensorFlow: ', tf.__version__)


vocab_size = 10000
LSTM_DIM = 256
BATCH_SIZE = 16
embedding_matrix = np.random.randn(vocab_size, 300)

class Encoder(tf.keras.Model):
    def __init__(self, vocab_size, embedding_matrix, n_units, batch_size):
        super(Encoder, self).__init__()
        self.n_units = n_units
        self.batch_size = batch_size

        self.embedding = Embedding(vocab_size, embedding_matrix.shape[1], weights=[embedding_matrix], trainable=True, mask_zero=True)
        self.lstm = LSTM(n_units, return_sequences=True, return_state=True, recurrent_initializer="glorot_uniform")


    @tf.function
    def call(self, inputs):
        input_utterence, initial_state = inputs
        input_embed = self.embedding(input_utterence)
        encoder_states, h1, c1 = self.lstm(input_embed, initial_state=initial_state)
        return encoder_states, h1, c1


    def create_initial_state(self):
        return tf.zeros((self.batch_size, self.n_units))

random_input = tf.random.uniform(shape=[BATCH_SIZE, 3], maxval=vocab_size, dtype=tf.int32)


encoder = Encoder(vocab_size, embedding_matrix, LSTM_DIM, BATCH_SIZE)

initial_state = [encoder.create_initial_state(), encoder.create_initial_state()]
_ = encoder([random_input, initial_state]) # required so that encoder.build is triggered


tf.saved_model.save(encoder, "encoder_model", signatures=encoder.call.get_concrete_function(
         [
             tf.TensorSpec(shape=[None, None], dtype=tf.int32, name='input_utterence'), 
             [
                 tf.TensorSpec(shape=[None, LSTM_DIM], dtype=tf.float32, name='initial_h'), 
                 tf.TensorSpec(shape=[None, LSTM_DIM], dtype=tf.float32, name='initial_c')
             ]
         ]))

loaded_model = tf.saved_model.load('encoder_model')
loaded_model([random_input, initial_state])

output:

TensorFlow:  2.2.0-rc1
WARNING:tensorflow:From /home/dl_user/tf_stable/lib/python3.7/site-packages/tensorflow/python/ops/resource_variable_ops.py:1817: calling BaseResourceVariable.__init__ (from tensorflow.python.ops.resource_variable_ops) with constraint is deprecated and will be removed in a future version.
Instructions for updating:
If using Keras pass *_constraint arguments to layers.
INFO:tensorflow:Assets written to: encoder_model/assets
(,
 ,
 )

Tensorflow saving subclass model which has multiple arguments to call() method

1 Answers