How to copy a tf.keras.models.Model subclass?

Question

I need to copy a keras model and there is no way that I know of which can be done unless the model is not a tf.keras.models.Model() subclass.

Note: The use copy.deepcopy() will work without giving any errors however it will result in another error whenever the copy is used.

Example:

import tensorflow as tf

class MyModel(tf.keras.Model):

  def __init__(self):
    super(MyModel, self).__init__()
    self.dense1 = tf.keras.layers.Dense(4, activation=tf.nn.relu)
    self.dense2 = tf.keras.layers.Dense(5, activation=tf.nn.softmax)
    self.dropout = tf.keras.layers.Dropout(0.5)

  def call(self, inputs, training=False):
    x = self.dense1(inputs)
    if training:
      x = self.dropout(x, training=training)
    return self.dense2(x)


if __name__ == '__main__':
    model1 = MyModel()
    model2 = tf.keras.models.clone_model(model1)

Results in:

Traceback (most recent call last):
  File "/Users/emadboctor/Library/Application Support/JetBrains/PyCharm2020.3/scratches/scratch.py", line 600, in <module>
    model2 = tf.keras.models.clone_model(model1)
  File "/usr/local/lib/python3.8/site-packages/tensorflow/python/keras/models.py", line 430, in clone_model
    return _clone_functional_model(
  File "/usr/local/lib/python3.8/site-packages/tensorflow/python/keras/models.py", line 171, in _clone_functional_model
    raise ValueError('Expected `model` argument '
ValueError: Expected `model` argument to be a functional `Model` instance, but got a subclass model instead.

Answer 1

Currently, we can't use tf.keras.models.clone_model for subclassed model API whereas we can for sequential and functional API. From doc,

model   Instance of Model (could be a functional model or a Sequential model).

Here is a workaround for your need. It makes sense if we need to copy a trained model, where we can get some optimized parameters. So, the main task is we need to create a new model by copying an existing model. The most convenient way for now of this scenario is to get trained weight and set to the newly created model instances. Let first build a model, train it and then get and set weight matrices to the new model.

import tensorflow as tf
import numpy as np

class ModelSubClassing(tf.keras.Model):
    def __init__(self, num_classes):
        super(ModelSubClassing, self).__init__()
        self.conv1 = tf.keras.layers.Conv2D(32, 3, strides=2, activation="relu")
        self.gap   = tf.keras.layers.GlobalAveragePooling2D()
        self.dense = tf.keras.layers.Dense(num_classes)

    def call(self, input_tensor, training=False):
        # forward pass: block 1 
        x = self.conv1(input_tensor)
        x = self.gap(x)
        return self.dense(x)

    def build_graph(self, raw_shape):
        x = tf.keras.layers.Input(shape=raw_shape)
        return tf.keras.Model(inputs=[x], outputs=self.call(x))


# compile 
sub_classing_model = ModelSubClassing(10)
sub_classing_model.compile(
        loss      = tf.keras.losses.CategoricalCrossentropy(),
        metrics   = tf.keras.metrics.CategoricalAccuracy(),
        optimizer = tf.keras.optimizers.Adam())

# plot for debug 
tf.keras.utils.plot_model(
    sub_classing_model.build_graph(x_train.shape[1:]),
    show_shapes=False,
    show_dtype=False,
    show_layer_names=True,
    expand_nested=False,
    dpi=96,
)

DataSet

(x_train, y_train), (_, _) = tf.keras.datasets.mnist.load_data()

# train set / data 
x_train = np.expand_dims(x_train, axis=-1)
x_train = x_train.astype('float32') / 255
# train set / target 
y_train = tf.keras.utils.to_categorical(y_train, num_classes=10)

# fit 
sub_classing_model.fit(x_train, y_train, batch_size=128, epochs=1)
# 469/469 [==============================] - 2s 2ms/step - loss: 8.2821

New Model / Copy

For the subclasses model, we have to initiate the class object.

sub_classing_model_copy = ModelSubClassing(10)
sub_classing_model_copy.build((x_train.shape))
sub_classing_model_copy.set_weights(sub_classing_model.get_weights()) # <- get and set wg

# plot for debug ; same as original plot
# but know, layer name is no longer same
# i.e. if, old: conv2d_40 , new/copy: conv2d_41
tf.keras.utils.plot_model(
    sub_classing_model_copy.build_graph(x_train.shape[1:]),
    show_shapes=False,
    show_dtype=False,
    show_layer_names=True,
    expand_nested=False,
    dpi=96,
)

Answer 2

def clones(module, N):
  
    Creation of N identical layers.
    :param module: module to clone
    :param N: number of copies
    :return: keras model of module copies
  
    seqm=KM.Sequential()
    for i in range(N):
        m = copy.deepcopy(module)
        m.name=m.name+str(i)
        seqm.add(m)
    return seqm