trainloop 0.8.0__tar.gz → 0.10.0__tar.gz

{trainloop-0.8.0 → trainloop-0.10.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: trainloop
-Version: 0.8.0
+Version: 0.10.0
 Summary: Minimal PyTorch training loop with hooks and checkpointing.
 Author: Karim Knaebel
 Author-email: Karim Knaebel <contact@knaebel.dev>

{trainloop-0.8.0 → trainloop-0.10.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "trainloop"
-version = "0.8.0"
+version = "0.10.0"
 description = "Minimal PyTorch training loop with hooks and checkpointing."
 readme = "README.md"
 requires-python = ">=3.10"
@@ -21,6 +21,7 @@ dev = [
     "zensical>=0.0.10",
     "mkdocstrings-python>=2.0.1",
     "pillow>=11.3.0",
+    "tensorboard>=2.20.0",
     "wandb>=0.20.1",
     "numpy>=2.2.6",
 ]

{trainloop-0.8.0 → trainloop-0.10.0}/pyproject.toml.orig

@@ -1,6 +1,6 @@
 [project]
 name = "trainloop"
-version = "0.8.0"
+version = "0.10.0"
 description = "Minimal PyTorch training loop with hooks and checkpointing."
 readme = "README.md"
 authors = [
@@ -22,6 +22,7 @@ dev = [
     "zensical>=0.0.10",
     "mkdocstrings-python>=2.0.1",
     "pillow>=11.3.0",
+    "tensorboard>=2.20.0",
     "wandb>=0.20.1",
     "numpy>=2.2.6",
 ]

{trainloop-0.8.0 → trainloop-0.10.0}/src/trainloop/__init__.py

@@ -6,6 +6,7 @@ from .hooks import (
     ImageFileLoggerHook,
     ProgressHook,
     StatsHook,
+    TensorBoardHook,
     TrainingStats,
     WandbHook,
 )
@@ -21,6 +22,7 @@ __all__ = [
     "LossNoneWarning",
     "ProgressHook",
     "StatsHook",
+    "TensorBoardHook",
     "TrainingStats",
     "WandbHook",
     "map_nested_tensor",

{trainloop-0.8.0 → trainloop-0.10.0}/src/trainloop/hooks.py

@@ -26,11 +26,18 @@ except ImportError:
     # only needed for WandbHook
     pass
 
+try:
+    import numpy as np
+    from torch.utils.tensorboard import SummaryWriter
+except ImportError:
+    # only needed for TensorBoardHook
+    pass
+
 try:
     from PIL import Image
     from PIL.Image import Image as PILImage
 except ImportError:
-    # only needed for WandbHook JPEG conversion
+    # only needed for WandbHook and TensorBoardHook image conversion
     pass
 
 from .trainer import BaseTrainer, Records
@@ -76,7 +83,9 @@ class BaseHook:
     def on_after_train(self, trainer: BaseTrainer):
         pass
 
-    def on_log(self, trainer: BaseTrainer, records: dict, dry_run: bool = False):
+    def on_log_scalars(
+        self, trainer: BaseTrainer, records: dict, dry_run: bool = False
+    ):
         pass
 
     def on_log_images(self, trainer: BaseTrainer, records: dict, dry_run: bool = False):
@@ -620,6 +629,9 @@ class EMAHook(BaseHook):
 class WandbHook(BaseHook):
     """Log metrics and images to Weights & Biases (rank 0 only).
 
+    Nested scalar key components are joined with ``/``. Avoid ``/`` within an
+    individual component because it is also W&B's panel namespace separator.
+
     Args:
         project: W&B project name.
         config: Optional config dict or JSON file path to log.
@@ -671,7 +683,9 @@ class WandbHook(BaseHook):
         if _dist_rank() == 0:
             self.wandb.finish()
 
-    def on_log(self, trainer: BaseTrainer, records: dict, dry_run: bool = False):
+    def on_log_scalars(
+        self, trainer: BaseTrainer, records: dict, dry_run: bool = False
+    ):
         if _dist_rank() == 0:
             data = {"/".join(k): v for k, v in flatten_nested_dict(records).items()}
             if not dry_run:
@@ -738,6 +752,109 @@ class WandbHook(BaseHook):
         return None
 
 
+class TensorBoardHook(BaseHook):
+    """Log metrics and images to TensorBoard (rank 0 only).
+
+    Nested namespace components are joined with ``namespace_separator`` and
+    the final name is separated with ``/``. Avoid these separators within
+    individual keys because TensorBoard uses them to organize tags.
+
+    Args:
+        texts: Optional text values to log when training starts.
+        namespace_separator: Separator for nested tag prefixes; the final TensorBoard
+            name boundary is always "/".
+    """
+
+    def __init__(
+        self,
+        texts: dict[str, Any] | None = None,
+        namespace_separator: str = "/",
+    ):
+        self.writer = None
+        self.purge_step = None
+        self.texts = texts
+        self.namespace_separator = namespace_separator
+
+    def on_before_train(self, trainer: BaseTrainer):
+        if _dist_rank() == 0:
+            self._open_writer(trainer)
+            for k, v in (self.texts or {}).items():
+                self.writer.add_text(k, str(v))
+            self.writer.flush()
+
+    def on_load_state_dict(self, trainer: BaseTrainer, state_dict: dict):
+        self.purge_step = trainer.step + 1  # the step that will be logged next
+        if self.writer is not None:
+            self.writer.close()
+            self._open_writer(trainer)
+
+    def on_after_train(self, trainer: BaseTrainer):
+        if self.writer is None:
+            return
+
+        self.writer.close()
+        self.writer = None
+
+    def on_log_scalars(
+        self, trainer: BaseTrainer, records: dict, dry_run: bool = False
+    ):
+        if self.writer is None:
+            return
+
+        data = {self._format_tag(k): v for k, v in flatten_nested_dict(records).items()}
+        if not dry_run:
+            for k, v in data.items():
+                self.writer.add_scalar(k, v, global_step=trainer.step)
+            self.writer.flush()
+        else:
+            trainer.logger.debug(f"Dry run log. Would log: {data}")
+
+    def on_log_images(self, trainer: BaseTrainer, records: dict, dry_run: bool = False):
+        if self.writer is None:
+            return
+
+        data = {
+            self._format_tag((*k, "image")): np.asarray(self._pil_to_rgb(v))
+            for k, v in flatten_nested_dict(records).items()
+        }
+        if not dry_run:
+            for k, v in data.items():
+                self.writer.add_image(k, v, global_step=trainer.step, dataformats="HWC")
+            self.writer.flush()
+        else:
+            trainer.logger.debug(f"Dry run log. Would log: {data}")
+
+    def _open_writer(self, trainer: BaseTrainer):
+        self.writer = SummaryWriter(
+            log_dir=str(trainer.workspace),
+            purge_step=self.purge_step,
+            max_queue=1_000,  # we already flush after every step
+        )
+        self.purge_step = None
+
+    def _format_tag(self, key: tuple[str, ...]):
+        *namespace, name = key
+        return (
+            f"{self.namespace_separator.join(namespace)}/{name}" if namespace else name
+        )
+
+    @staticmethod
+    def _pil_to_rgb(img: "PILImage", bg_color: tuple = (255, 255, 255)):
+        if img.mode == "RGB":
+            return img
+        elif img.mode == "L":
+            return img.convert("RGB")
+        elif img.mode in ("RGBA", "LA"):
+            background = Image.new("RGB", img.size, bg_color)
+            background.paste(img, mask=img.getchannel("A"))
+            return background
+        else:
+            warnings.warn(
+                f"Trying to convert {img.mode} to RGB in a best-effort manner."
+            )
+            return img.convert("RGB")
+
+
 class ImageFileLoggerHook(BaseHook):
     """Persist logged images to ``workspace/visualizations`` on rank 0.
 

{trainloop-0.8.0 → trainloop-0.10.0}/src/trainloop/trainer.py

@@ -349,17 +349,22 @@ class BaseTrainer:
     def unwrapped_model(self):
         return self.unwrap(self.model)
 
-    def log(self, records: dict[str, Any], dry_run: bool = False):
+    def log_scalars(self, records: dict[str, Any], dry_run: bool = False):
         """
-        Dispatch numeric records to hooks (e.g., trackers or stdout).
+        Dispatch nested scalar records to logging hooks.
+
+        Dictionary keys are treated as path components. Logging hooks may join
+        nested components with backend-specific namespace separators, such as
+        ``/``. Avoid those separators within individual keys because distinct
+        nested paths may otherwise map to the same backend metric name.
 
         Args:
-            records: Nested dict of numeric metrics to log.
+            records: Nested dict of scalar metrics to log.
             dry_run: If True, hooks should avoid side effects and only report intent.
         """
-        self.logger.debug("log()")
+        self.logger.debug("log_scalars()")
         for h in self.hooks:
-            h.on_log(self, records, dry_run=dry_run)
+            h.on_log_scalars(self, records, dry_run=dry_run)
 
     def log_images(self, records: dict[str, Any], dry_run: bool = False):
         """