trainloop 0.5.0__tar.gz → 0.5.1__tar.gz

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

@@ -1,9 +1,9 @@
 Metadata-Version: 2.3
 Name: trainloop
-Version: 0.5.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.5.0 → trainloop-0.5.1}/pyproject.toml

@@ -1,10 +1,10 @@
 [project]
 name = "trainloop"
-version = "0.5.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.5.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},
                 }
             }
         )