Converting TensorFlow Keras model API to model subclassing

For a simple TF2 Object detection CNN architecture defined using Keras’s functional API, a batch of data is obtained as:

 example, label = next(data_generator(batch_size = 32)) example.keys() # dict_keys(['image']) image = example['image'] image.shape # (32, 144, 144, 3) label.keys() # dict_keys(['class_out', 'box_out']) label['class_out'].shape, label['box_out'].shape # ((32, 9), (32, 2)) 

The CNN architecture defined using Keras’s functional API is:

 input_ = Input(shape = (144, 144, 3), name = 'image') # name - An optional name string for the Input layer. Should be unique in # a model (do not reuse the same name twice). It will be autogenerated if it isn't provided. # Here 'image' is the Python3 dict's key used to map the data to one of the layer in the model. x = input_ # Define a conv block- x = Conv2D(filters = 64, kernel_size = 3, activation = 'relu')(x) x = BatchNormalization()(x) x = MaxPool2D(pool_size = 2)(x) x = Flatten()(x) # flatten the last pooling layer's output volume x = Dense(256, activation='relu')(x) # We are using a data generator which yields dictionaries. Using 'name' argument makes it # possible to map the correct data generator's output to the appropriate layer class_out = Dense(units = 9, activation = 'softmax', name = 'class_out')(x) # classification output box_out = Dense(units = 2, activation = 'linear', name = 'box_out')(x) # regression output # Define the CNN model- model = tf.keras.models.Model(input_, [class_out, box_out]) # since we have 2 outputs, we use a list 

I am attempting to define it using Model sub-classing as:

 class OD(Model): def __init__(self): super(OD, self).__init__() self.conv1 = Conv2D(filters = 64, kernel_size = 3, activation = None) self.bn = BatchNormalization() self.pool = MaxPool2D(pool_size = 2) self.flatten = Flatten() self.dense = Dense(256, activation = None) self.class_out = Dense(units = 9, activation = None, name = 'class_out') self.box_out = Dense(units = 2, activation = 'linear', name = 'box_out') def call(self, x): x = tf.nn.relu(self.bn(self.conv1(x))) x = self.pool(x) x = self.flatten(x) x = tf.nn.relu(self.dense(x)) x = [tf.nn.softmax(self.class_out(x)), self.box_out(x)] return x A batch of training data is obtained as: example, label = next(data_generator(batch_size = 32)) example.keys() # dict_keys(['image']) image = example['image'] image.shape # (32, 144, 144, 3) label.keys() # dict_keys(['class_out', 'box_out']) label['class_out'].shape, label['box_out'].shape # ((32, 9), (32, 2)) 

Is my Model sub-classing architecture equivalent to Keras’s functional API?