tensorflow: `tf.keras.Model.save` does not support subclassed model when saving model as SavedModel format

I’ve been running OP’s snippet and the bug is clearly reproducible. As @jvishnuvardhan pointed, using _set_inputs does the work. However, I think it should not be necessary if the model is called and fed an input. In other words, the call function of Model should automatically call _set_inputs with the input received

@gugarosa One workaround is that after finished custom training, call model.predict on some input tensors and then call model.save to save your model as SavedModel format.

@jvishnuvardhan So do you mean that model.save can not be used if I’m writing custom training loop instead of compilation + fit? I should call tf.saved_model.save(model), is it right?

I am having the same problem. I have defined my subclass custom model and its custom training loop. Furthermore, I am also not using .compile() nor .fit(), and I can’t try to figure which saving method I could use.

I have tested every single of them and could not find any way to solve it. Additionally, TF does not provide any documentation regarding saving a model with custom training (without using .compile() and .fit()).

I followed this guide and it worked for me: https://www.tensorflow.org/guide/saved_model#exporting_custom_models

In other words, the below code should work. That said, I do agree that an explicit example for exporting a sub-classed keras model with a custom training loop is missing on the tf page.

import tensorflow as tf


class Model(tf.keras.Model):
    def __init__(self):
        super(Model, self).__init__()
        self.d = tf.keras.layers.Dense(2)

    @tf.function(input_signature=[tf.TensorSpec([2, 3], tf.float32)])
    def call(self, x, training=True, mask=None):
        return self.d(x)


model = Model()
model(tf.random.normal((2, 3)))

tf.saved_model.save(model, '/tmp/serving_path')

I can’t save a custom model even though I successfully call model.predict(…) before the save. I get all kinds of error. My model is a seq2seq model. Why is tensorflow saying the subclassing keras.Model is the recommended low-level/flexible way when i is almost impossible to save it as a SavedModel?..

hi @jvishnuvardhan ,

As you didn’t run the your model

subclassed model needs to be called on some data in order to create its weights.

But I indeed called my model, see

model = Model()
# next line !
model(tf.random.normal((2, 3)))

Please reopen the issue as it is actually a bug.

@zakizhou The error ValueError: Model <__main__.Model object at 0x7f7a6bad1e48> cannot be saved because the input shapes have not been set. Usually, input shapes are automatically determined from calling .fit() or .predict(). To manually set the shapes, call model._set_inputs(inputs). clearly mention that input_shape is missing.

Please follow the steps mentioned in the TF webiste for Saving Subclassed Models. It clearly mentions that

First of all, a subclassed model that has never been used cannot be saved.

That's because a subclassed model needs to be called on some data in order to create its weights.

Until the model has been called, it does not know the shape and dtype of the input data it should be expecting, and thus cannot create its weight variables. You may remember that in the Functional model from the first section, the shape and dtype of the inputs was specified in advance (via keras.Input(...)) -- that's why Functional models have a state as soon as they're instantiated.

As you didn’t run the your model, it does not know the shape and dtype of the input data it should be expecting, and thus cannot create its weight variables.

I am closing the issue. Please feel free to open if you face any error after implementing your model following the steps outlined in TF website. There are couple of tutorials and guidelines on TF websites that will be helpful to you. Thanks!

@andreclaudino @aingo03304 @AlexFuster Can you please open a new issue with details of your model and share a simple standalone code to reproduce the issue. Thanks!