trainloop 0.4.0__tar.gz → 0.5.1__tar.gz

{trainloop-0.4.0 → trainloop-0.5.1}/PKG-INFO

@@ -1,9 +1,9 @@
 Metadata-Version: 2.3
 Name: trainloop
-Version: 0.4.0
+Version: 0.5.1
 Summary: Minimal PyTorch training loop with hooks and checkpointing.
-Author: Karim Abou Zeid
-Author-email: Karim Abou Zeid <contact@ka.codes>
+Author: Karim Knaebel
+Author-email: Karim Knaebel <contact@knaebel.dev>
 Requires-Dist: pillow>=11.3.0
 Requires-Dist: torch>=2.0.0
 Requires-Python: >=3.10

{trainloop-0.4.0 → trainloop-0.5.1}/pyproject.toml

@@ -1,10 +1,10 @@
 [project]
 name = "trainloop"
-version = "0.4.0"
+version = "0.5.1"
 description = "Minimal PyTorch training loop with hooks and checkpointing."
 readme = "README.md"
 authors = [
-    { name = "Karim Abou Zeid", email = "contact@ka.codes" }
+    { name = "Karim Knaebel", email = "contact@knaebel.dev" }
 ]
 requires-python = ">=3.10"
 dependencies = [

{trainloop-0.4.0 → trainloop-0.5.1}/src/trainloop/hooks.py

@@ -269,7 +269,7 @@ class ProgressHook(_StatsHook):
     ):
         eta = self.eta_tracker.get_eta(trainer.max_steps - trainer.step)
         trainer.logger.info(
-            f"Step {trainer.step:>{len(str(trainer.max_steps))}}/{trainer.max_steps}:"
+            f"Step {trainer.step}/{trainer.max_steps}:"
             + f" step {step_time:.4f}{'s' if self.show_units else ''} data {data_time:.4f}{'s' if self.show_units else ''}"
             + (f" eta {eta}" if eta is not None else "")
             + (
@@ -309,6 +309,13 @@ class LoggingHook(_StatsHook):
     ):
         super().__init__(interval, sync)
 
+    def on_before_step(self, trainer: BaseTrainer):
+        super().on_before_step(trainer)
+        self.lrs = [
+            (param_group.get("name", f"group_{i}"), param_group["lr"])
+            for i, param_group in enumerate(trainer.optimizer.param_groups)
+        ]  # record the LR before the scheduler steps
+
     def process_stats(
         self,
         trainer: BaseTrainer,
@@ -320,10 +327,6 @@ class LoggingHook(_StatsHook):
         max_memory: float | None,
         records: Records,
     ):
-        lrs = [
-            (param_group.get("name", f"group_{i}"), param_group["lr"])
-            for i, param_group in enumerate(trainer.optimizer.param_groups)
-        ]
         trainer.log(
             {
                 "train": records
@@ -334,7 +337,7 @@ class LoggingHook(_StatsHook):
                     "data_time": data_time,
                     "step_time": step_time,
                     "non_finite_grad_retry_count": non_finite_grad_retry_count,
-                    "lr": {name: lr for name, lr in lrs},
+                    "lr": {name: lr for name, lr in self.lrs},
                 }
             }
         )
@@ -346,7 +349,8 @@ class CheckpointingHook(BaseHook):
     Args:
         interval: Save every ``interval`` steps.
         keep_previous: Keep the last N checkpoints in addition to the latest.
-        keep_interval: Keep checkpoints every ``keep_interval`` steps.
+        keep_interval: Save and keep checkpoints every ``keep_interval`` steps.
+        keep_steps: Save and keep checkpoints at these explicit step numbers.
         path: Directory (relative to workspace unless absolute) for checkpoints.
         load: Path to load at startup or ``\"latest\"`` to auto-resume.
         exit_signals: Signals that trigger a checkpoint then exit.
@@ -358,7 +362,9 @@ class CheckpointingHook(BaseHook):
         self,
         interval: int,
         keep_previous: int = 0,  # keep N previous checkpoints
-        keep_interval: int = 0,  # keep checkpoints of every N-th step
+        keep_interval: int = 0,  # save and keep checkpoints of every N-th step
+        keep_steps: Sequence[int]
+        | None = None,  # save and keep checkpoints at these steps
         path: Path | str = "checkpoints",
         load: Path | str | Literal["latest"] | None = "latest",
         exit_signals: list[signal.Signals] | signal.Signals = None,
@@ -370,6 +376,7 @@ class CheckpointingHook(BaseHook):
         self.interval = interval
         self.keep_previous = keep_previous
         self.keep_interval = keep_interval
+        self.keep_steps = set(keep_steps or [])
         self.path = Path(path)
         self.load_path = Path(load) if load is not None else None
 
@@ -435,18 +442,23 @@ class CheckpointingHook(BaseHook):
             save_and_exit = exit_signal != -1
 
         # NOTE: Check if last step here (not in on_after_train) to avoid saving twice
-        if (
+        should_keep = trainer.step in self.keep_steps or (
+            self.keep_interval > 0 and trainer.step % self.keep_interval == 0
+        )
+        should_save = (
             trainer.step % self.interval == 0
             or trainer.step == trainer.max_steps
+            or should_keep
             or save_and_exit
-        ):
+        )
+        if should_save:
             if save_and_exit:
                 trainer.logger.info(
                     f"=> Caught signal {exit_signal}. Saving checkpoint before exit ..."
                 )
             self._save_checkpoint(
                 trainer,
-                keep=self.keep_interval > 0 and trainer.step % self.keep_interval == 0,
+                keep=should_keep,
             )
             if save_and_exit:
                 _dist_barrier()

{trainloop-0.4.0 → trainloop-0.5.1}/src/trainloop/trainer.py

@@ -276,7 +276,9 @@ class BaseTrainer:
                             < self.max_non_finite_grad_retries
                         ):
                             non_finite_grad_retry_count += 1
-                            self.step_info["non_finite_grad_retry_count"] = non_finite_grad_retry_count
+                            self.step_info["non_finite_grad_retry_count"] = (
+                                non_finite_grad_retry_count
+                            )
                             self.logger.warning(
                                 f"Gradient is non-finite. Retrying step {self.step} (retry {non_finite_grad_retry_count}"
                                 + (