Comparison of defining models

models.py

import tensorflow as tf

# Reset states (to be sure)
tf.keras.backend.clear_session()

# Method 1
x = tf.keras.Input(shape=(224, 224, 3), dtype=tf.uint8)
c = tf.keras.layers.Flatten()(x)
y = tf.keras.layers.Dense(units=2, activation="softmax")(c)
model1 = tf.keras.Model(inputs=x, outputs=y)

# Reset states (to be sure)
tf.keras.backend.clear_session()

# Method 1
class MyModel(tf.keras.Model):
    def __init__(self, input_shape, num_classes, **kwargs):
        
        # Init
        super(MyModel, self).__init__(**kwargs)
        
        # Define input
        self.inp = tf.keras.layers.Input(input_shape, dtype=tf.uint8)
        
        # Define layers
        self.flatten = tf.keras.layers.Flatten()
        self.fc = tf.keras.layers.Dense(units=num_classes, activation="softmax")      
       
        # Define output
        self.out = self.call(self.inp)
        
        # Reinit 
        super(MyModel, self).__init__(inputs=self.inp, outputs=self.out, **kwargs)
            
    def call(self, x, training=False):
        c = self.flatten(x)
        c = self.fc(c)
        return c
    
    def build(self):
        self._is_graph_network = True
        self._init_graph_network(inputs=self.inp, outputs=self.out)
        
model2 = MyModel(input_shape=(1, 224, 224, 3), num_classes=2)

# Compare layers
for l in model1.layers:
    print(l)    
for l in model2.layers:
    print(l)

# Compare summaries
model1.summary() 
model2.summary()

# Compare plotting (execute in seperate notebook cells to see the plot)
tf.keras.utils.plot_model(model1, to_file="/tmp/model.png", show_shapes=True, show_layer_names=True)
tf.keras.utils.plot_model(model2, to_file="/tmp/model.png", show_shapes=True, show_layer_names=True)

# Compare saving
tf.keras.models.save_model(model1, "/tmp/test_model_saving")
tf.keras.models.save_model(model2, "/tmp/test_model_saving")