[Python] Model Summary not picking up from the model

I am fine tuning a CodeBert Model using my custom Dataset and Tokenizers. I tried to unfreeze 4 the last 4 layers of the model. When checking if the layers are unfreezed it shows me which layers are not frozen. Following is the code that I am using.

def create_model(
bert_model_name="microsoft/codebert-base",
num_classes=442,
dense_units=[628, 528],
dropout_rate=0.3,
learning_rate=1e-4,
trainable_bert=False,):
bert_model = RobertaModel.from_pretrained(bert_model_name)
bert_model.trainable = trainable_bert

for i, layer in enumerate(bert_model.roberta.encoder.layer):
if i < len(bert_model.roberta.encoder.layer) - 4:
layer.trainable = False

if layer.trainable:
print(layer.name, "is trainable")
input_ids = tf.keras.Input(shape=(512,), dtype=tf.int32, name="input_ids")
attention_mask = tf.keras.Input(shape=(512,), dtype=tf.int32, name="attention_mask")

def bert_layer(inputs):
last_hidden = bert_model(inputs).last_hidden_state
return last_hidden[:, 0, :]

x = tf.keras.layers.Lambda(
bert_layer,
output_shape=(768,)
)({"input_ids": input_ids, "attention_mask": attention_mask})

for units in dense_units:
x = tf.keras.layers.Dense(units)(x)
x = tf.keras.layers.BatchNormalization()(x)
x = tf.keras.layers.Activation("relu")(x)
x = tf.keras.layers.Dropout(dropout_rate)(x)

output = tf.keras.layers.Dense(num_classes, activation="linear")(x)

model = tf.keras.Model(inputs=[input_ids, attention_mask], outputs=output)

model.compile(
optimizer=tf.keras.optimizers.Adam(learning_rate=learning_rate),
loss=tf.keras.losses.BinaryCrossentropy(from_logits=True),
metrics=[lambda y_true, y_pred: positive_class_accuracy(y_true, y_pred, k=5, prob_threshold=0.01)]
)

return model


This is now printing the following where it shows that the last 4 layers are unfrozen

layer_._8 is trainable
layer_._9 is trainable
layer_._10 is trainable
layer_._11 is trainable


But my model.summary() does not show the last 4 layers My confusion is that the CodeBert Model is not being showed in the model.summary() at all and this is not training my model.

Continue reading...