feat: Initial Commit

utils/checkpoint.py

@@ -0,0 +1,85 @@
+import os
+import torch
+
+
+def save_checkpoint(
+        model,
+        optimizer,
+        epoch,
+        loss,
+        path,
+):
+    """
+    Save a training checkpoint.
+    """
+
+    os.makedirs(os.path.dirname(path), exist_ok=True)
+
+    torch.save(
+        {
+            "epoch": epoch,
+            "model_state_dict": model.state_dict(),
+            "optimizer_state_dict": optimizer.state_dict(),
+            "loss": loss,
+        },
+        path,
+    )
+
+
+def load_checkpoint(
+        path,
+        model,
+        optimizer=None,
+        map_location="cpu",
+):
+    """
+    Load a training checkpoint.
+
+    Returns:
+        epoch (int)
+        loss (float)
+    """
+
+    checkpoint = torch.load(
+        path,
+        map_location=map_location,
+    )
+
+    model.load_state_dict(
+        checkpoint["model_state_dict"]
+    )
+
+    if optimizer is not None:
+        optimizer.load_state_dict(
+            checkpoint["optimizer_state_dict"]
+        )
+
+    return (
+        checkpoint["epoch"],
+        checkpoint["loss"],
+    )
+
+
+def latest_checkpoint(folder):
+    """
+    Returns the newest checkpoint file in a folder.
+    """
+
+    if not os.path.exists(folder):
+        return None
+
+    checkpoints = [
+        os.path.join(folder, file)
+        for file in os.listdir(folder)
+        if file.endswith(".pth")
+    ]
+
+    if len(checkpoints) == 0:
+        return None
+
+    checkpoints.sort(
+        key=os.path.getmtime,
+        reverse=True,
+    )
+
+    return checkpoints[0]