Subclassing tf.keras.models.Model save_model saves h5py but not h5/hdf5 file types

[Ryandry1st]

Please make sure that this is a bug. As per our GitHub Policy, we only address code/doc bugs, performance issues, feature requests and build/installation issues on GitHub. tag:bug_template

System information

• Have I written custom code (as opposed to using a stock example script provided in TensorFlow): Yes
• OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Google Colab
• TensorFlow installed from (source or binary): Pip
• TensorFlow version (use command below): 2.0.0b
• Python version: 3.6.7
• GPU model and memory: Tesla T4

Describe the current behavior
-Using tf 2.0.0b-gpu on google colab.

While using the subclassing API for a subclassed layer and model, I was unable to use the model.save_model() function for h5 or hdf5 file types, but I could successfully save and load the model if it was saved as h5py file type. In the toy example being used it worked correctly, although this may not be the case. Note that the get_config method was implemented in the custom layer and custom model.

Describe the expected behavior
Either the save_model should always work (I believe this is a feature goal) and the documentation should reflect this, or if the save is likely to produce incorrect results it should raise an error and the documentation should continue to suggest that custom models can only be saved with the save_weights feature.

Code to reproduce the issue

import tensorflow as tf
from tensorflow import keras

class resblock(keras.layers.Layer):
    def __init__(self, n_layers, n_neurons, **kwargs):
        super().__init__(**kwargs)
        self.hidden = [keras.layers.Dense(n_neurons,
                                          activation='elu',
                                          kernel_initializer='he_normal')
                       for _ in range(n_layers)]
        
    def call(self, inputs):
        z = inputs
        for layer in self.hidden:
            z = layer(z)
        return inputs + z
    
    def get_config(self):
        base_config = super().get_config()
        return {**base_config}

class res_mod(keras.models.Model):
    def __init__(self, output_dim, activation=None, **kwargs):
        super().__init__(**kwargs)
        self.f1 = keras.layers.Flatten()
        self.hidden1 = keras.layers.Dense(100, activation='elu', kernel_initializer='he_normal')
        self.b1 = resblock(2, 100)
        self.b2 = resblock(2, 100)
        self.output1 = keras.layers.Dense(output_dim, activation=keras.activations.get(activation))
        
    def call(self, inputs):
        z = self.f1(inputs)
        z = self.hidden1(z)
        for _ in range(4):
            z = self.b1(z)
        z = self.b2(z)
        return self.output1(z)
    
    def get_config(self):
        base_config = super().get_config()
        return{**base_config, "output_dim" : output_dim, "activation": activation}
    

model = res_mod(10, activation='softmax')
model.compile(loss='sparse_categorical_crossentropy', optimizer='nadam', metrics=['accuracy'])

model.fit(train, epochs=25, validation_data=test)

# This is able to save and works correctly, returning the trained model
model.save('custom_model.h5py')
del model
model = keras.models.load_model('custom_model.h5py', custom_objects={'resblock': resblock})

Other info / logs
This will raise an error that only sequential or functional models can be saved
model.save('custom_model.h5')

[gadagashwini-zz]

@Ryandry1st I tried reproducing the issue but looks code snippet is incomplete. Please provide minimal code snippet to reproduce issue reported here. Also provide the error message that you have got. Thanks!

[Ryandry1st]

The only code not provided was the actual data for the train and test data objects. It is possible to reproduce with the MNIST/ fashion MNIST or really any other dataset, simply change the 10 from the line
model = res_mod(10, activation='softmax')
to whatever you are using it for.
Assuming you have loaded train and test data then the code executes up to the save_model function, the error is produced by attempting to save the model as either h5 or hdf5 (while h5py works fine) and produces the following error:

NotImplementedError: Saving the model to HDF5 format requires the model to be a Functional model or a Sequential model. It does not work for subclassed models, because such models are defined via the body of a Python method, which isn't safely serializable. Consider saving to the Tensorflow SavedModel format (by setting save_format="tf") or using save_weights.

The entire traceback is attached in the text file.

Custom_Model_Save_Error.txt

[Ryandry1st]

The point of this issue is that if it cannot be safely serialized then saving as an h5py either is handling it improperly, and should raise an error similar to the above, OR it works as it should with h5py (based on only this one test case and limited testing besides checking weights and accuracy) and the documentation should include this information, and probably adopt the method/recommend it.

[Ryandry1st]

Any luck on this issue? One option for the data is of course just randomly generating x and y values. It does not matter what data is fed into the network, only that it is a trained subclassed network that is attempted to save

[ymodak]

The code snippet you have provided looks incomplete. Can you please check with latest tf-2.0--nightly version?
pip install tf-2.0-nightly-preview

[Ryandry1st]

Here is a more basic code snippet that throws the error. I do want to stress though that the only difference is that I am providing some randomly generated x and y data, and removing layers from the custom model for simplicity. There is no real new code or any ideas and it was no stretch to reach the warning using the previous code, and using random data as was suggested a month ago.

x = tf.random.uniform((100,))
y = tf.random.uniform((100,))


class test_model(keras.models.Model):
    def __init__(self, output_dim, activation=None, **kwargs):
        super().__init__(**kwargs)
        self.f1 = keras.layers.Flatten()
        self.hidden1 = keras.layers.Dense(100, activation='elu', kernel_initializer='he_normal', name="h1")
        self.output1 = keras.layers.Dense(output_dim, activation=keras.activations.get(activation))
        
    def call(self, inputs):
        z = self.f1(inputs)
        z = self.hidden1(z)
        return self.output1(z)
    
    def get_config(self):
        base_config = super().get_config()
        return{**base_config, "output_dim" : output_dim, "activation": activation}
    
model = test_model(1)
model.compile(loss='mse', optimizer='adam', metrics=['mae'])
model.fit(x, y, epochs=5)
print(model.weights[0])
model.save('custom_model.hdf5')
del model
model = keras.models.load_model('custom_model.hdf5')
print(model.weights[0])

You will see that this throws the error, but the following would not.

x = tf.random.uniform((100,))
y = tf.random.uniform((100,))


class test_model(keras.models.Model):
    def __init__(self, output_dim, activation=None, **kwargs):
        super().__init__(**kwargs)
        self.f1 = keras.layers.Flatten()
        self.hidden1 = keras.layers.Dense(100, activation='elu', kernel_initializer='he_normal', name="h1")
        self.output1 = keras.layers.Dense(output_dim, activation=keras.activations.get(activation))
        
    def call(self, inputs):
        z = self.f1(inputs)
        z = self.hidden1(z)
        return self.output1(z)
    
    def get_config(self):
        base_config = super().get_config()
        return{**base_config, "output_dim" : output_dim, "activation": activation}
    
model = test_model(1)
model.compile(loss='mse', optimizer='adam', metrics=['mae'])
model.fit(x, y, epochs=5)
print(model.weights[0])
model.save('custom_model.h5py')
del model
model = keras.models.load_model('custom_model.h5py')
print(model.weights[0])

So somewhere along the way it should either be adopted, the error removed, or it should throw a warning using this file format as well.

[Ryandry1st]

I also tried using the preview with the more complicated code and found that it throws a new error, which appears to be due to improperly recalling custom objects. This may be some of the premise behind the warning, but it did work when using the beta version and no longer works with the nightly preview.

/usr/local/lib/python3.6/dist-packages/tensorflow_core/python/eager/function.py in _convert_inputs_to_signature(inputs, input_signature, flat_input_signature)
   1701       flatten_inputs)):
   1702     raise ValueError("Python inputs incompatible with input_signature:\n%s" %
-> 1703                      format_error_message(inputs, input_signature))
   1704 
   1705   if need_packing:

ValueError: Python inputs incompatible with input_signature:
  inputs: (
    Tensor("res_mod_1_2/Cast:0", shape=(None, 16), dtype=float32))
  input_signature: (
    TensorSpec(shape=(None, 16), dtype=tf.float64, name='input_1'))

[k-w-w]

In TensorFlow 2.0, you should save the model with save_format="h5" extension, otherwise it will default to saving to the SavedModel format. Another way to save to H5 is to set the extension to ".h5" or ".hdf5" or ".keras". The H5 format does not support saving Subclassed models, so the initial error stating that only Sequential and Functional models may be saved is intended.

(The documentation describes the argument on this page: https://www.tensorflow.org/versions/r2.0/api_docs/python/tf/keras/Model#save)

---

I have an internal change that should fix the ValueError (Python inputs incompatible with input_signature), but I can't be sure because the examples above do not raise the error. I'll update this post when the change is copied to github.