TF2.6: ValueError: Model cannot be saved because the input shapes have not been set

Question:

I want to create a custom model using transfer learning in Google Colab.

import tensorflow as tf
from tensorflow.keras.layers import Conv2D
from tensorflow.python.keras.applications.xception import Xception

class MyModel(tf.keras.Model):

  def __init__(self, input_shape, num_classes=5, dropout_rate=0.5):
    super(MyModel, self).__init__()
    self.weight_dict = {}
    self.weight_dict['backbone'] = Xception(input_shape=input_shape, weights='imagenet', include_top=False)
    self.weight_dict['outputs'] = Conv2D(num_classes, (1, 1), padding="same", activation="softmax")
    self.build((None,) + input_shape)

  def call(self, inputs, training=False):
    self.weight_dict['backbone'].trainable = False
    x = self.weight_dict['backbone'](inputs)
    x = self.weight_dict['outputs'](x)
    return x

model = MyModel(input_shape=(256, 256, 3))
model.save('./saved')

However, I encounter this error:

ValueError: Model `<__main__.MyModel object at 0x7fc66134bdd0>` 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.build(input_shape)`.

Yes, there is no call to .fit() or .predict(). But there is a call to .build in the __init__() method of the class. What am I to do?

Asked By: Jim

||

Answers:

If the layer has not been built, compute_output_shape will call build on the layer. This assumes that the layer will later be used with inputs that match the input shape provided.

Working code as shown below

import tensorflow as tf
print(tf.__version__)
from tensorflow.keras.layers import Conv2D
from tensorflow.keras.applications.xception import Xception

class MyModel(tf.keras.Model):

  def __init__(self, input_shape, num_classes=5, dropout_rate=0.5):
    super(MyModel, self).__init__()
    self.weight_dict = {}
    self.weight_dict['backbone'] = Xception(input_shape=input_shape, weights='imagenet', include_top=False)
    self.weight_dict['outputs'] = Conv2D(num_classes, (1, 1), padding="same", activation="softmax")
    self.build((None,) + input_shape)

  def call(self, inputs, training=False):
    self.weight_dict['backbone'].trainable = False
    x = self.weight_dict['backbone'](inputs)
    x = self.weight_dict['outputs'](x)
    return x

input_shape=(256, 256, 3)
model=MyModel(input_shape)

model.compute_output_shape(input_shape=(None, 256, 256, 3))
model.save('./saved')

Output:

2.6.0

Downloading data from https://storage.googleapis.com/tensorflow/keras-applications/xception/xception_weights_tf_dim_ordering_tf_kernels_notop.h5
83689472/83683744 [==============================] - 1s 0us/step
INFO:tensorflow:Assets written to: ./saved/assets

For more information you can refer here.

Answered By: TFer2

If you build e.g. a GNN with multiple inputs of variable size, the proposal of TFer2 won’t work. Specifying the TensorSpecs in the decorator AND using tf.saved_model.save instead works:

import keras.layers
import tensorflow as tf
import unittest


class TestModel(keras.Model):
    def __init__(self):
        super(TestModel, self).__init__()
        self.w = tf.Variable(initial_value=tf.initializers.he_normal()(shape=[10, 10]))

    @tf.function(input_signature=[tf.TensorSpec([10, None], tf.float32), tf.TensorSpec([10, None], tf.float32)])  # This line defines the inputs' sizes of the network call
    def __call__(self, x, y):
        return tf.matmul(self.w, x) + tf.matmul(self.w, y)


class SaveAndLoadTest(unittest.TestCase):
    def __init__(self):
        super(SaveAndLoadTest, self).__init__()

        x = tf.ones([10, 5])
        y = tf.ones([10, 5])

        model = TestModel()
        z = model(x, y)

        tf.saved_model.save(model, "/tmp/test_model/")  # saving a this way works
        with self.assertRaises(ValueError):
            model.save("/tmp/test_model/")  # saving a model this way fails, regardless of assigning TensorSpecs to tf.function

        model_loaded = tf.saved_model.load("/tmp/test_model/")
        z_loaded = model_loaded(x, y)

        self.assertTrue((z_loaded.numpy() == z.numpy()).all())  # making sure the outputs are the same


if __name__ == "__main__":
    SaveAndLoadTest()
    print("Success.")
Answered By: FlorianJWirth