Allow training to restart from a folder (#64)

* If directory passed, get last model

* Save the final model after all iterations

Author: eigenvivek

src/xvr/cli/commands/restart.py

@@ -32,12 +32,19 @@ def restart(
     Restart model training from a checkpoint.
     """
     import os
+    from pathlib import Path
 
     import torch
     import wandb
 
     from ...model import Trainer
 
+    # If ckptpath is a directory, get the last saved model
+    ckptpath = Path(ckptpath)
+    if ckptpath.is_dir():
+        ckptpath = sorted(ckptpath.glob("*.pth"))[-1]
+    ckptpath = str(ckptpath)
+
     # Load the config from the previous model checkpoint
     config = torch.load(ckptpath, weights_only=False)["config"]
     config["ckptpath"] = ckptpath

src/xvr/cli/commands/train.py

@@ -339,6 +339,12 @@ def train(
     # Create the output directory for saving model weights
     Path(outpath).mkdir(parents=True, exist_ok=True)
 
+    # If ckptpath is a directory, get the last saved model
+    ckptpath = Path(ckptpath)
+    if ckptpath.is_dir():
+        ckptpath = sorted(ckptpath.glob("*.pth"))[-1]
+    ckptpath = str(ckptpath)
+
     # Parse 6-DoF pose parameters
     alphamin, alphamax = r1
     betamin, betamax = r2

src/xvr/model/trainer.py

@@ -168,6 +168,9 @@ def train(self, run=None):
             if run is not None:
                 self._log_wandb(itr, log, imgs, masks)
 
+        # Save the final model
+        self._checkpoint(itr)
+
     def step(self, itr):
         if self.single_subject:
             log, imgs, masks = self._step_single_subject(itr)