[Python] KFold cross-validation in Keras: model not resetting between folds (MobileNet...

I am trying to perform KFold cross-validation on a Keras model. The first fold runs exactly as expected, but from the second fold onwards the model doesn’t seem to reset. The training behaves strangely, as if the weights from the previous fold are being reused instead of starting fresh. Below is a snippet of the logs

Epoch 14/15
[00:12<00:00, 2.45s/it]
pose_loss: 0.1407 - quality_loss: 1.1151
pose_accuracy: 1.0000 - quality_accuracy: 0.3625
pose_precision: 1.0000 - pose_recall: 1.0000
quality_precision: 0.2493 - quality_recall: 0.3581
val_pose_loss: 0.1631 - val_quality_loss: 2.1690
val_pose_accuracy: 0.9062 - val_quality_accuracy: 0.1562
val_pose_precision: 0.5417 - val_pose_recall: 0.4792
val_quality_precision: 0.2024 - val_quality_recall: 0.2000
Epoch 14: Last 25 backbone layers unfrozen with L2 reg (1e-06)
Recompiling models due to regularization changes...
Epoch 15/15
100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 5/5 [00:12<00:00, 2.45s/it]
pose_loss: 0.1310 - quality_loss: 1.1538
pose_accuracy: 1.0000 - quality_accuracy: 0.3250
pose_precision: 1.0000 - pose_recall: 1.0000
quality_precision: 0.2124 - quality_recall: 0.3256
val_pose_loss: 0.1766 - val_quality_loss: 2.1545
val_pose_accuracy: 0.9062 - val_quality_accuracy: 0.1562
val_pose_precision: 0.5417 - val_pose_recall: 0.4792
val_quality_precision: 0.2024 - val_quality_recall: 0.2000
Fold 1 completed - Final val_pose_acc: 0.9062, Final val_quality_acc: 0.1562

=== FOLD 2/6 ===
Train samples: 81, Validation samples: 16
Backbone frozen: True
=== OPTIMIZER STATE AFTER RESET ===
Pose optimizer iterations: 0
Quality optimizer iterations: 0
Backbone optimizer iterations: 0
Percentage Difference between estimation weights: -1.1358185111887449
Percentage Difference between backbone weights: -0.9901197477751298
Model weights, learning rates, and optimizers have been reset for new fold.
Epoch 0: Backbone frozen
Epoch 1/15
100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 5/5 [00:12<00:00, 2.47s/it]
pose_loss: 0.1797 - quality_loss: 1.1455
pose_accuracy: 0.9750 - quality_accuracy: 0.3375
pose_precision: 0.9750 - pose_recall: 0.9806
quality_precision: 0.4066 - quality_recall: 0.3272
val_pose_loss: 1.8522 - val_quality_loss: 1.1337
val_pose_accuracy: 0.1875 - val_quality_accuracy: 0.3750
val_pose_precision: 0.0312 - val_pose_recall: 0.1667
val_quality_precision: 0.1250 - val_quality_recall: 0.3333


From my understanding, every fold should start with a newly initialized model. But in my case, it looks like the model is somehow being cached or reused by Keras/TensorFlow.

And I can assure that there is no data leakage in the pipeline.

The function used for model building and model resetting are mentioned below respectively.

def build_model(self):
mobilenet = None
self.backbone = None
estimation_model = None
rating_model = None

inputs = tf.keras.layers.Input(shape=self.input_shape)
mobilenet = tf.keras.applications.MobileNetV2(weights='imagenet', include_top=False, input_shape=(224, 224, 3))
mobilenet.trainable = False
self.backbone = mobilenet

base_cnn = tf.keras.layers.TimeDistributed(mobilenet)(inputs)


x = tf.keras.layers.Conv3D(128, (3, 3, 3), padding='same', activation='relu')(base_cnn)
x = tf.keras.layers.BatchNormalization()(x)
x = tf.keras.layers.MaxPooling3D((1, 2, 2))(x)

x = tf.keras.layers.Conv3D(64, (3, 3, 3), padding='same', activation='relu')(x)
x = tf.keras.layers.BatchNormalization()(x)

# Global pooling
x = tf.keras.layers.GlobalAveragePooling3D()(x)

x = tf.keras.layers.Dropout(0.35)(x)
shared = tf.keras.layers.Dense(128, activation='relu', kernel_regularizer=tf.keras.regularizers.l2(1e-2))(x)

backbone_model = tf.keras.Model(inputs=inputs, outputs=shared, name='Backbone')

pose_input = tf.keras.layers.Input(shape=(128,))
est = tf.keras.layers.Dropout(0.5)(pose_input)
pose_estimation = tf.keras.layers.Dense(self.num_poses, activation='softmax', name='pose')(est)

estimation_model = tf.keras.Model(inputs=pose_input, outputs=pose_estimation, name='estimation')

rating_input = tf.keras.layers.Input(shape=(128,))
rating = tf.keras.layers.Dense(32, activation='relu')(rating_input)
rating =tf.keras.layers.Dropout(0.5)(rating)
rating = tf.keras.layers.Dense(8, activation='relu')(rating)
pose_rating = tf.keras.layers.Dense(self.num_rating, activation='softmax')(rating)
rating_model = tf.keras.Model(inputs=rating_input, outputs=pose_rating, name='quality')

return backbone_model, estimation_model, rating_model


def _reset_model_weights(self):
"""Reset model weights to initial state for fresh training in each fold."""
backbone_model, estimation_model, rating_model = self.model
old_estimation_weights = estimation_model.layers[-1].kernel.numpy().copy() if hasattr(estimation_model.layers[-1], 'kernel') else None
old_backbone_weights = backbone_model.layers[-1].kernel.numpy().copy() if hasattr(backbone_model.layers[-1], 'kernel') else None

self.model = self.build_model()

# Reset learning rates
self.pose_lr = self.original_pose_lr
self.rating_lr = self.original_rating_lr
self.model_needs_recompile = False

backbone_model, estimation_model, rating_model = self.model

# Reset backbone reference
backbone_model, _, _ = self.model
for layer in backbone_model.layers:
if isinstance(layer, tf.keras.layers.TimeDistributed):
self.backbone = layer.layer
break

if self.backbone:
self.backbone.trainable = False
print(f"Backbone frozen: {not self.backbone.trainable}")

self.compile_models()


Versions: Tensorflow = 2.18.0, keras = 3.9.2

1. Is this behavior expected because of the way Keras handles models internally?
2. Or is there some additional step required to clear/reset the model between folds?
3. Could it be caused by the MobileNet backbone, since I am unfreezing it during training?

Continue reading...