Can't save custom subclassed model

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.

This will be fixed in an upcoming release according to the 1.13 pre-release patch notes:

• Keras & Python API:
  • Subclassed Keras models can now be saved through tf.contrib.saved_model.save_keras_model.

EDIT: It seems this is not quite as finished as the notes suggest. The docs for that function for v1.13 state:

Model limitations: - Sequential and functional models can always be saved. - Subclassed models can only be saved when serving_only=True. This is due to the current implementation copying the model in order to export the training and evaluation graphs. Because the topology of subclassed models cannot be determined, the subclassed models cannot be cloned. Subclassed models will be entirely exportable in the future.

Tensorflow 2.1 allows to save subclassed models with SavedModel format

From my beginning using Tensorflow, i was always a fan of Model Subclass, i feel this way of build models more pythonic and collaborative friendly. But saving the model was always a point of pain with this approach.

Recently i started to update my knowledge and reach to the following information that seems to be True for Tensorflow 2.1:

Subclassed Models

I found this

Second approach is by using model.save to save whole model and by using load_model to restore previously stored subclassed model.

This last saves the model, the weight and other stuff into a SavedModel file

And by last the confirmation:

Saving custom objects: If you are using the SavedModel format, you can skip this section. The key difference between HDF5 and SavedModel is that HDF5 uses object configs to save the model architecture, while SavedModel saves the execution graph. Thus, SavedModels are able to save custom objects like subclassed models and custom layers without requiring the orginal code.

I tested this personally, and efectively, model.save() for subclassed models generate a SavedModel save. There is no more need for use model.save_weights() or related functions, they now are more for specific usecases.

This is suposed to be the end of this painful path for all of us interested in Model Subclassing.