trainer: Enable resuming from checkpoint

Author: lisavader

trainer/train_model.py

@@ -91,6 +91,12 @@ def get_args():
         default=1,
         help='How many dataset chunks to pool, shuffle, and train on in a single fit call (default: 1)'
     )
+    parser.add_argument(
+        '--resume-from-checkpoint',
+        type=str,
+        default=None,
+        help='Path to .keras checkpoint file to resume training from (default: None)'
+    )
     args, _ = parser.parse_known_args()
     return args
 
@@ -278,20 +284,58 @@ def categorical_crossentropy_2d_gene_ml(y_true, y_pred):
     except Exception:
         pass
 
-    if N_GPUS > 1:
-        # https://keras.io/guides/distributed_training/
-        strategy = tensorflow.distribute.MirroredStrategy()
-        tee('Number of devices: {}'.format(strategy.num_replicas_in_sync))
-        with strategy.scope():
+    if args.resume_from_checkpoint:
+        tee(f"\033[1mLoading model from checkpoint: {args.resume_from_checkpoint}\033[0m")
+        # Load with custom objects
+        custom_objects = {
+            'categorical_crossentropy_2d_gene_ml': categorical_crossentropy_2d_gene_ml,
+            'WarmupSchedule': WarmupSchedule,
+            'DecayOnPlateauSchedule': DecayOnPlateauSchedule,
+        }
+        model = keras.models.load_model(args.resume_from_checkpoint, custom_objects=custom_objects)
+
+        # Extract the decay_factor from the saved optimizer config
+        # The learning rate schedule is saved in the optimizer, and we need to extract
+        # the decay_factor value to continue modifying it if needed
+        try:
+            lr_config = model.optimizer.learning_rate.get_config()
+            saved_decay_factor = lr_config.get('decay_factor', 1.0)
+            tee(f"Loaded decay_factor from checkpoint: {saved_decay_factor}")
+        except AttributeError:
+            saved_decay_factor = 1.0
+            tee("Could not extract decay_factor from checkpoint, assuming 1.0 (no decay applied yet)")
+
+        # Create a new decay_factor variable initialized to the saved value
+        decay_factor = tensorflow.Variable(saved_decay_factor, trainable=False, dtype=tensorflow.float32, name='lr_decay_factor')
+
+        # Get current learning rate
+        current_lr = model.optimizer.learning_rate
+        if callable(current_lr):
+            lr_value = float(tensorflow.keras.backend.get_value(current_lr(model.optimizer.iterations)))
+        else:
+            lr_value = float(tensorflow.keras.backend.get_value(current_lr))
+        tee(f"Resumed from checkpoint. Current learning rate: {lr_value:.5f}")
+
+        # Extract starting epoch from checkpoint filename (e.g., "ep10" -> start at 11)
+        match = re.search(r'_ep(\d+)\.keras', args.resume_from_checkpoint)
+        start_epoch = int(match.group(1)) + 1 if match else 1
+        tee(f"\033[1mResuming from epoch {start_epoch}\033[0m")
+    else:
+        if N_GPUS > 1:
+            # https://keras.io/guides/distributed_training/
+            strategy = tensorflow.distribute.MirroredStrategy()
+            tee('Number of devices: {}'.format(strategy.num_replicas_in_sync))
+            with strategy.scope():
+                model = GeneML(L, W, AR, num_classes)
+                model.compile(loss=loss,
+                              optimizer=keras.optimizers.Adam(learning_rate=lr_schedule,
+                                                              weight_decay=WEIGHT_DECAY))
+        else:
             model = GeneML(L, W, AR, num_classes)
             model.compile(loss=loss,
                           optimizer=keras.optimizers.Adam(learning_rate=lr_schedule,
                                                           weight_decay=WEIGHT_DECAY))
-    else:
-        model = GeneML(L, W, AR, num_classes)
-        model.compile(loss=loss,
-                      optimizer=keras.optimizers.Adam(learning_rate=lr_schedule,
-                                                      weight_decay=WEIGHT_DECAY))
+        start_epoch = 1
     # model.summary()
 
     ###############################################################################
@@ -374,7 +418,7 @@ def print_performance_metrics(indices, max_eval):
 
         return acceptor_score, donor_score
 
-    for epoch_num in range(1, args.num_epochs + 1):
+    for epoch_num in range(start_epoch, args.num_epochs + 1):
         # Shuffle indices for this epoch (no replacement)
         # Use epoch-based seed for deterministic but different permutation per epoch
         np.random.seed(SEED + epoch_num)