Making a custom keras architecture

model.py

import numpy as np 
import os
import numpy as np
import tensorflow as tf
from tensorflow.keras.models import *
from tensorflow.keras.layers import *

def model(pretrained_weights = None,input_size = (224,224,3), outchannel = 25, dropout_p = 0.25):
  inputTensor = tf.keras.Input(IMG_SHAPE)
  ## Feature extractor network
  ResNet152 = tf.keras.applications.ResNet152(include_top=False, weights='imagenet', input_tensor=inputTensor)
  # Extracting the desired layers. You can use multiple layer in case of SSD. 
  # Check here : https://github.com/rs9899/mySSDimplementation/blob/master/MobileNetSSD_v2.ipynb
  block = ResNet152.get_layer('conv5_block3_out').output
  ## You custom layers
  C1 = tf.keras.layers.Conv2DTranspose(
        256, [4,4], strides=2, padding='same', 
        dilation_rate=(1, 1), activation='relu', use_bias=True,
        kernel_initializer='he_normal', bias_initializer='zeros'
    )
  out1 = C1(block)
  out1 = tf.keras.layers.Dropout(dropOut_P)(out1)
  ### ..... many such of your layers
  Final = tf.keras.layers.Conv2D(outchannel, 1)
  finalOut = Final(out1)
  
  model = tf.keras.Model(inputs = inputTensor, outputs = finalOut)

  if(pretrained_weights):
      model.load_weights(pretrained_weights)

  return model

train.py

### THIS ONE IS HIGHLY INCOMPLETE

import tensorflow as tf
import numpy as np
from model import model as mymodel
import tensorflow_addons as tfa

class CustomModelCheckpoint(tf.keras.callbacks.Callback):
  def on_epoch_end(self, epoch, logs=None):
    self.model.save_weights('./weightPath.hdf5', overwrite=True)
    
if __name__ == "__main__":
  ### argparser
  args = argparse.ArgumentParser()
  args.add_argument('--trainset_path', type=str, default="./traindata")
  config = args.parse_args()
  
  model = mymodel()
  opt = tfa.optimizers.RectifiedAdam(lr=config.learning_rate)
  model.compile(loss= tf.keras.losses.BinaryCrossentropy() , optimizer=opt , metrics = ['accuracy'])
  cbk = CustomModelCheckpoint()
#   train_gen = Data generators
#   STEP_SIZE_TRAIN = trainDataSize // BatchSize
#   model.fit( x=train_gen, epochs = NumEpoch,
#               steps_per_epoch = STEP_SIZE_TRAIN,
#               callbacks = [cbk])
  print("Training completed")