From 4c4d6713e06203d36fb0ab1109dfefd35024535d Mon Sep 17 00:00:00 2001
From: YannAhlgrim <yann@ahlgrim.net>
Date: Wed, 15 Apr 2026 13:43:05 +0200
Subject: [PATCH] add save checkpoint logic

---
 src/train_supervised.py | 101 ++++++++++++++++++++++++++++++----------
 1 file changed, 77 insertions(+), 24 deletions(-)

diff --git a/src/train_supervised.py b/src/train_supervised.py
index 383971a..8821668 100644
--- a/src/train_supervised.py
+++ b/src/train_supervised.py
@@ -39,6 +39,16 @@ def main(args):
     m_args = args["meta"]
     o_args = args["optimization"]
     d_args = args["data"]
+    l_args = args["logging"]
+
+    # -- Paths for saving
+    folder = l_args["folder"]
+    tag = l_args["write_tag"]
+    if not os.path.exists(folder):
+        os.makedirs(folder, exist_ok=True)
+
+    save_path = os.path.join(folder, f"{tag}" + "-ep{epoch}.pth.tar")
+    latest_path = os.path.join(folder, f"{tag}-latest.pth.tar")
 
     # -- 1. Initialize Encoder
     encoder, _ = init_model(
@@ -52,7 +62,7 @@ def main(args):
     embed_dim = m_args.get("embed_dim")
     model = ViTClassifier(encoder, m_args["num_classes"], embed_dim).to(device)
 
-    # -- 3. Handle Freezing (Linear Probing vs Fine-Tuning)
+    # -- 3. Handle Freezing
     if o_args["freeze_weights"]:
         logger.info("Freezing encoder weights (Linear Probing mode)")
         for name, param in model.encoder.named_parameters():
@@ -60,26 +70,7 @@ def main(args):
     else:
         logger.info("Training full model (Fine-tuning mode)")
 
-    # -- 4. Load Pre-trained Weights
-    if m_args["load_checkpoint"]:
-        load_path = os.path.join(args["logging"]["folder"], m_args["read_checkpoint"])
-        checkpoint = torch.load(load_path, map_location="cpu")
-        msg = model.encoder.load_state_dict(checkpoint["encoder"], strict=False)
-        logger.info(f"Loaded encoder from {load_path} with msg: {msg}")
-
-    # -- 5. Data Setup
-    transform = make_transforms(crop_size=d_args["crop_size"])
-    _, loader, sampler = make_iwildcam(
-        transform=transform,
-        split="train",
-        batch_size=d_args["batch_size"],
-        root_path=d_args["root_path"],
-        rank=rank,
-        world_size=world_size,
-        collator=None,  # No mask collator needed
-    )
-
-    # -- 6. Optimizer Selection
+    # -- 4. Optimizer Selection
     params = [p for p in model.parameters() if p.requires_grad]
     if o_args["optimizer"].lower() == "adamw":
         optimizer = torch.optim.AdamW(
@@ -90,11 +81,70 @@ def main(args):
             params, lr=o_args["lr"], momentum=0.9, weight_decay=o_args["weight_decay"]
         )
 
+    # -- 5. Resume/Load Logic
+    start_epoch = 0
+    # Priority 1: Check if we are resuming from an interrupted run (latest-path)
+    # Priority 2: Check if we are loading a specific pre-trained checkpoint (m_args["load_checkpoint"])
+
+    checkpoint_to_load = None
+    resuming_interrupted = False
+
+    if os.path.exists(latest_path):
+        checkpoint_to_load = latest_path
+        resuming_interrupted = True
+    elif m_args["load_checkpoint"]:
+        checkpoint_to_load = os.path.join(folder, m_args["read_checkpoint"])
+
+    if checkpoint_to_load:
+        checkpoint = torch.load(checkpoint_to_load, map_location="cpu")
+
+        if resuming_interrupted:
+            # Load full state to resume exactly where we left off
+            model.load_state_dict(checkpoint["model"])
+            optimizer.load_state_dict(checkpoint["opt"])
+            start_epoch = checkpoint["epoch"]
+            logger.info(
+                f"Resuming training from {checkpoint_to_load} at epoch {start_epoch}"
+            )
+        else:
+            # Loading just encoder weights for a fresh supervised run
+            msg = model.encoder.load_state_dict(checkpoint["encoder"], strict=False)
+            logger.info(
+                f"Loaded pre-trained encoder from {checkpoint_to_load} with msg: {msg}"
+            )
+
+    # -- 6. Data Setup
+    transform = make_transforms(crop_size=d_args["crop_size"])
+    _, loader, sampler = make_iwildcam(
+        transform=transform,
+        split="train",
+        batch_size=d_args["batch_size"],
+        root_path=d_args["root_path"],
+        rank=rank,
+        world_size=world_size,
+        collator=None,
+    )
+
     criterion = nn.CrossEntropyLoss().to(device)
     model = DistributedDataParallel(model, device_ids=[torch.cuda.current_device()])
 
-    # -- 7. Training Loop
-    for epoch in range(o_args["epochs"]):
+    # -- 7. Define Save Function
+    def save_checkpoint(epoch, current_loss):
+        save_dict = {
+            "model": model.module.state_dict(),  # model.module because of DDP
+            "opt": optimizer.state_dict(),
+            "epoch": epoch,
+            "loss": current_loss,
+            "args": args,
+        }
+        if rank == 0:
+            torch.save(save_dict, latest_path)
+            if epoch % checkpoint_freq == 0:
+                torch.save(save_dict, save_path.format(epoch=epoch))
+            logger.info(f"Checkpoint saved at epoch {epoch}")
+
+    # -- 8. Training Loop
+    for epoch in range(start_epoch, o_args["epochs"]):
         sampler.set_epoch(epoch)
         model.train()
         loss_meter = AverageMeter()
@@ -116,9 +166,12 @@ def main(args):
 
             if itr % 10 == 0 and rank == 0:
                 logger.info(
-                    f"Epoch {epoch} [{itr}/{len(loader)}] Loss: {loss_meter.avg:.4f}"
+                    f"Epoch {epoch + 1} [{itr}/{len(loader)}] Loss: {loss_meter.avg:.4f}"
                 )
 
+        # Save at the end of every epoch
+        save_checkpoint(epoch + 1, loss_meter.avg)
+
 
 if __name__ == "__main__":
     main()