python 3.x - Can't save custom subclassed model

Question

Inspired by tf.keras.Model subclassing I created custom model.
I can train it and get successfull results, but I can't save it.
I use python3.6 with tensorflow v1.10 (or v1.9)

Minimal complete code example here:

import tensorflow as tf
from tensorflow.keras.datasets import mnist


class Classifier(tf.keras.Model):
    def __init__(self):
        super().__init__(name="custom_model")

        self.batch_norm1 = tf.layers.BatchNormalization()
        self.conv1 = tf.layers.Conv2D(32, (7, 7))
        self.pool1 = tf.layers.MaxPooling2D((2, 2), (2, 2))

        self.batch_norm2 = tf.layers.BatchNormalization()
        self.conv2 = tf.layers.Conv2D(64, (5, 5))
        self.pool2 = tf.layers.MaxPooling2D((2, 2), (2, 2))

    def call(self, inputs, training=None, mask=None):
        x = self.batch_norm1(inputs)
        x = self.conv1(x)
        x = tf.nn.relu(x)
        x = self.pool1(x)

        x = self.batch_norm2(x)
        x = self.conv2(x)
        x = tf.nn.relu(x)
        x = self.pool2(x)

        return x


if __name__ == '__main__':
    (x_train, y_train), (x_test, y_test) = mnist.load_data()

    x_train = x_train.reshape(*x_train.shape, 1)[:1000]
    y_train = y_train.reshape(*y_train.shape, 1)[:1000]

    x_test = x_test.reshape(*x_test.shape, 1)
    y_test = y_test.reshape(*y_test.shape, 1)

    y_train = tf.keras.utils.to_categorical(y_train)
    y_test = tf.keras.utils.to_categorical(y_test)

    model = Classifier()

    inputs = tf.keras.Input((28, 28, 1))

    x = model(inputs)
    x = tf.keras.layers.Flatten()(x)
    x = tf.keras.layers.Dense(10, activation="sigmoid")(x)

    model = tf.keras.Model(inputs=inputs, outputs=x)
    model.compile(optimizer="adam", loss="binary_crossentropy", metrics=["accuracy"])
    model.fit(x_train, y_train, epochs=1, shuffle=True)

    model.save("./my_model")

Error message:

1000/1000 [==============================] - 1s 1ms/step - loss: 4.6037 - acc: 0.7025
Traceback (most recent call last):
  File "/home/user/Data/test/python/mnist/mnist_run.py", line 62, in <module>
    model.save("./my_model")
  File "/home/user/miniconda3/envs/ml3.6/lib/python3.6/site-packages/tensorflow/python/keras/engine/network.py", line 1278, in save
    save_model(self, filepath, overwrite, include_optimizer)
  File "/home/user/miniconda3/envs/ml3.6/lib/python3.6/site-packages/tensorflow/python/keras/engine/saving.py", line 101, in save_model
    'config': model.get_config()
  File "/home/user/miniconda3/envs/ml3.6/lib/python3.6/site-packages/tensorflow/python/keras/engine/network.py", line 1049, in get_config
    layer_config = layer.get_config()
  File "/home/user/miniconda3/envs/ml3.6/lib/python3.6/site-packages/tensorflow/python/keras/engine/network.py", line 1028, in get_config
    raise NotImplementedError
NotImplementedError

Process finished with exit code 1

I looked into the error line and found out that get_config method checks self._is_graph_network

Do anybody deal with this problem?

Thanks!

Update 1:
On the keras 2.2.2 (not tf.keras)
Found comment (for model saving)
file: keras/engine/network.py
Function: get_config

# Subclassed networks are not serializable
# (unless serialization is implemented by
# the author of the subclassed network).

So, obviously it won't work...
I wonder, why don't they point it out in the documentation (Like: "Use subclassing without ability to save!")

Update 2:
Found in keras documentation:

In subclassed models, the model's topology is defined as Python code
(rather than as a static graph of layers). That means the model's
topology cannot be inspected or serialized. As a result, the following
methods and attributes are not available for subclassed models:

model.inputs and model.outputs.
model.to_yaml() and model.to_json()
model.get_config() and model.save().

So, there is no way to save model by using subclassing.
It's possible to only use Model.save_weights()

Answer 1

TensorFlow 2.2

Thanks for @cal for noticing me that the new TensorFlow has supported saving the custom models!

By using model.save to save the whole model and by using load_model to restore previously stored subclassed model. The following code snippets describe how to implement them.

class ThreeLayerMLP(keras.Model):

  def __init__(self, name=None):
    super(ThreeLayerMLP, self).__init__(name=name)
    self.dense_1 = layers.Dense(64, activation='relu', name='dense_1')
    self.dense_2 = layers.Dense(64, activation='relu', name='dense_2')
    self.pred_layer = layers.Dense(10, name='predictions')

  def call(self, inputs):
    x = self.dense_1(inputs)
    x = self.dense_2(x)
    return self.pred_layer(x)

def get_model():
  return ThreeLayerMLP(name='3_layer_mlp')

model = get_model()
# Save the model
model.save('path_to_my_model',save_format='tf')

# Recreate the exact same model purely from the file
new_model = keras.models.load_model('path_to_my_model')

See: Save and serialize models with Keras - Part II: Saving and Loading of Subclassed Models

TensorFlow 2.0

TL;DR:

1. do not use model.save() for custom subclass keras model;
2. use save_weights() and load_weights() instead.

---

With the help of the Tensorflow Team, it turns out the best practice of saving a Custom Sub-Class Keras Model is to save its weights and load it back when needed.

The reason that we can not simply save a Keras custom subclass model is that it contains custom codes, which can not be serialized safely. However, the weights can be saved/loaded when we have the same model structure and custom codes without any problem.

There has a great tutorial written by Francois Chollet who is the author of Keras, for how to save/load Sequential/Functional/Keras/Custom Sub-Class Models in Tensorflow 2.0 in Colab at here. In Saving Subclassed Models section, it said that:

Sequential models and Functional models are datastructures that represent a DAG of layers. As such, they can be safely serialized and deserialized.

A subclassed model differs in that it's not a datastructure, it's a piece of code. The architecture of the model is defined via the body of the call method. This means that the architecture of the model cannot be safely serialized. To load a model, you'll need to have access to the code that created it (the code of the model subclass). Alternatively, you could be serializing this code as bytecode (e.g. via pickling), but that's unsafe and generally not portable.