trainloop 0.7.0__tar.gz → 0.9.0__tar.gz

{trainloop-0.7.0 → trainloop-0.9.0}/PKG-INFO

@@ -1,11 +1,10 @@
 Metadata-Version: 2.3
 Name: trainloop
-Version: 0.7.0
+Version: 0.9.0
 Summary: Minimal PyTorch training loop with hooks and checkpointing.
 Author: Karim Knaebel
 Author-email: Karim Knaebel <contact@knaebel.dev>
-Requires-Dist: pillow>=11.3.0
-Requires-Dist: torch>=2.0.0
+Requires-Dist: torch>=2.1.0
 Requires-Python: >=3.10
 Description-Content-Type: text/markdown
 

trainloop-0.9.0/pyproject.toml

@@ -0,0 +1,27 @@
+[project]
+name = "trainloop"
+version = "0.9.0"
+description = "Minimal PyTorch training loop with hooks and checkpointing."
+readme = "README.md"
+requires-python = ">=3.10"
+dependencies = ["torch>=2.1.0"]
+
+[[project.authors]]
+name = "Karim Knaebel"
+email = "contact@knaebel.dev"
+
+[build-system]
+requires = ["uv_build>=0.9.0,<0.10.0"]
+build-backend = "uv_build"
+
+[dependency-groups]
+dev = [
+    "pytest>=8.4.0",
+    "ruff>=0.11.13",
+    "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.7.0/pyproject.toml → trainloop-0.9.0/pyproject.toml.orig

@@ -1,6 +1,6 @@
 [project]
 name = "trainloop"
-version = "0.7.0"
+version = "0.9.0"
 description = "Minimal PyTorch training loop with hooks and checkpointing."
 readme = "README.md"
 authors = [
@@ -8,8 +8,7 @@ authors = [
 ]
 requires-python = ">=3.10"
 dependencies = [
-    "pillow>=11.3.0",
-    "torch>=2.0.0",
+    "torch>=2.1.0",
 ]
 
 [build-system]
@@ -22,6 +21,8 @@ dev = [
     "ruff>=0.11.13",
     "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.7.0 → trainloop-0.9.0}/src/trainloop/__init__.py

@@ -1,27 +1,29 @@
 from .hooks import (
     BaseHook,
     CheckpointingHook,
-    CudaMaxMemoryHook,
-    EmaHook,
+    CUDAMaxMemoryHook,
+    EMAHook,
     ImageFileLoggerHook,
     ProgressHook,
     StatsHook,
+    TensorBoardHook,
     TrainingStats,
     WandbHook,
 )
 from .trainer import BaseTrainer, LossNoneWarning, map_nested_tensor
 
 __all__ = [
-    "BaseTrainer",
     "BaseHook",
+    "BaseTrainer",
+    "CUDAMaxMemoryHook",
     "CheckpointingHook",
-    "CudaMaxMemoryHook",
+    "EMAHook",
+    "ImageFileLoggerHook",
+    "LossNoneWarning",
     "ProgressHook",
     "StatsHook",
+    "TensorBoardHook",
     "TrainingStats",
-    "EmaHook",
     "WandbHook",
-    "ImageFileLoggerHook",
-    "LossNoneWarning",
     "map_nested_tensor",
 ]

{trainloop-0.7.0 → trainloop-0.9.0}/src/trainloop/hooks.py

@@ -5,21 +5,20 @@ import sys
 import tempfile
 import time
 import warnings
+from collections.abc import Callable, Iterable, Sequence
 from dataclasses import dataclass
 from datetime import timedelta
 from numbers import Number
 from pathlib import Path
-from typing import Any, Callable, Iterable, Literal, Sequence
+from typing import Any, Literal
 
 import torch
 import torch.distributed as dist
-from PIL import Image
-from PIL.Image import Image as PILImage
 from torch.distributed.checkpoint.state_dict import (
     get_model_state_dict,
     set_model_state_dict,
 )
-from torch.optim.swa_utils import AveragedModel, get_ema_avg_fn
+from torch.optim.swa_utils import AveragedModel, get_ema_multi_avg_fn
 
 try:
     import wandb
@@ -27,8 +26,26 @@ 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 and TensorBoardHook image conversion
+    pass
+
 from .trainer import BaseTrainer, Records
-from .utils import flatten_nested_dict, key_average
+from .utils import (
+    flatten_nested_dict,
+    key_average,
+    log_state_dict_incompatible_keys,
+)
 
 
 def _dist_is_initialized() -> bool:
@@ -85,8 +102,8 @@ class TrainingStats:
     grad_norm: float | None
     step_time: float
     data_time: float
-    non_finite_grad_retry_count: float
-    max_memory: float | None
+    non_finite_grad_retries: float
+    cuda_max_memory: float | None
     records: Records
     param_groups: list[dict[str, Any]]
 
@@ -119,8 +136,8 @@ class StatsHook(BaseHook):
         self.grad_norms = []
         self.data_times = []
         self.step_times = []
-        self.non_finite_grad_retry_counts = []
-        self.max_memories = []
+        self.non_finite_grad_retries = []
+        self.cuda_max_memories = []
 
     def on_before_step(self, trainer: BaseTrainer):
         super().on_before_step(trainer)
@@ -137,11 +154,11 @@ class StatsHook(BaseHook):
         self.records_ls.append(key_average(trainer.step_info["records"]))
         self.data_times.append(sum(trainer.step_info["data_time"]))  # total
         self.step_times.append(trainer.step_info["step_time"])
-        self.non_finite_grad_retry_counts.append(
-            trainer.step_info["non_finite_grad_retry_count"]
+        self.non_finite_grad_retries.append(
+            trainer.step_info["non_finite_grad_retries"]
         )
-        if "max_memory" in trainer.step_info:
-            self.max_memories.append(trainer.step_info["max_memory"])
+        if "cuda_max_memory" in trainer.step_info:
+            self.cuda_max_memories.append(trainer.step_info["cuda_max_memory"])
 
         if trainer.step % self.interval == 0 or trainer.step == trainer.max_steps:
             # aggregate over steps
@@ -150,10 +167,12 @@ class StatsHook(BaseHook):
             records = key_average(self.records_ls)
             data_time = sum(self.data_times) / len(self.data_times)
             step_time = sum(self.step_times) / len(self.step_times)
-            non_finite_grad_retry_count = sum(self.non_finite_grad_retry_counts) / len(
-                self.non_finite_grad_retry_counts
+            non_finite_grad_retries = sum(self.non_finite_grad_retries) / len(
+                self.non_finite_grad_retries
+            )
+            cuda_max_memory = (
+                max(self.cuda_max_memories) if self.cuda_max_memories else None
             )
-            max_memory = max(self.max_memories) if self.max_memories else None
 
             if self.sync and _dist_world_size() > 1:
                 # aggregate accross all ranks
@@ -168,18 +187,18 @@ class StatsHook(BaseHook):
                         "records": records,
                         "data_time": data_time,
                         "step_time": step_time,
-                        "non_finite_grad_retry_count": non_finite_grad_retry_count,
-                        "max_memory": max_memory,
+                        "non_finite_grad_retries": non_finite_grad_retries,
+                        "cuda_max_memory": cuda_max_memory,
                     },
                 )
                 records = key_average([stat["records"] for stat in gathered])
                 data_time = sum(stat["data_time"] for stat in gathered) / len(gathered)
                 step_time = sum(stat["step_time"] for stat in gathered) / len(gathered)
-                non_finite_grad_retry_count = sum(
-                    stat["non_finite_grad_retry_count"] for stat in gathered
+                non_finite_grad_retries = sum(
+                    stat["non_finite_grad_retries"] for stat in gathered
                 ) / len(gathered)
-                if "max_memory" in trainer.step_info:
-                    max_memory = max(stat["max_memory"] for stat in gathered)
+                if "cuda_max_memory" in trainer.step_info:
+                    cuda_max_memory = max(stat["cuda_max_memory"] for stat in gathered)
 
             self.callback(
                 trainer,
@@ -188,8 +207,8 @@ class StatsHook(BaseHook):
                     grad_norm=grad_norm.item() if grad_norm is not None else None,
                     step_time=step_time,
                     data_time=data_time,
-                    non_finite_grad_retry_count=non_finite_grad_retry_count,
-                    max_memory=max_memory,
+                    non_finite_grad_retries=non_finite_grad_retries,
+                    cuda_max_memory=cuda_max_memory,
                     records=records,
                     param_groups=self.param_groups,
                 ),
@@ -277,8 +296,8 @@ class ProgressHook(StatsHook):
             + f" step {stats.step_time:.4f}{'s' if self.show_units else ''} data {stats.data_time:.4f}{'s' if self.show_units else ''}"
             + (f" eta {eta}" if eta is not None else "")
             + (
-                f" mem {stats.max_memory:#.3g}{'GiB' if self.show_units else ''}"
-                if stats.max_memory is not None
+                f" mem {stats.cuda_max_memory:#.3g}{'GiB' if self.show_units else ''}"
+                if stats.cuda_max_memory is not None
                 else ""
             )
             + f" loss {stats.loss:.4f}"
@@ -333,7 +352,7 @@ class CheckpointingHook(BaseHook):
         | 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,
+        exit_signals: list[signal.Signals] | signal.Signals | None = None,
         exit_code: int | Literal["128+signal"] = "128+signal",
         exit_wait: timedelta | float = 0.0,
     ):
@@ -551,39 +570,54 @@ class CheckpointingHook(BaseHook):
             return False
 
 
-class CudaMaxMemoryHook(BaseHook):
+class CUDAMaxMemoryHook(BaseHook):
     """Record peak CUDA memory per step into ``trainer.step_info``."""
 
     def on_before_step(self, trainer: BaseTrainer):
         torch.cuda.reset_peak_memory_stats(trainer.device)
 
     def on_after_step(self, trainer: BaseTrainer):
-        trainer.step_info["max_memory"] = torch.cuda.max_memory_allocated(
+        trainer.step_info["cuda_max_memory"] = torch.cuda.max_memory_allocated(
             trainer.device
         ) / (1024**3)  # GiB
 
 
-class EmaHook(BaseHook):
+class EMAHook(BaseHook):
     """Maintain an exponential moving average of model weights.
 
     Args:
         decay: EMA decay rate.
+        use_buffers: Whether to include model buffers in the EMA.
     """
 
-    def __init__(self, decay: float):
+    def __init__(self, decay: float = 0.999, use_buffers: bool = False):
         self.decay = decay
+        self.use_buffers = use_buffers
 
     def on_before_train(self, trainer: BaseTrainer):
         trainer.logger.info("=> Creating EMA model ...")
         # Note that AveragedModel does not seem to support FSDP. It will crash here.
-        self.ema_model = AveragedModel(trainer.model, avg_fn=get_ema_avg_fn(self.decay))
+        self.ema_model = AveragedModel(
+            trainer.model,
+            multi_avg_fn=get_ema_multi_avg_fn(self.decay),
+            use_buffers=self.use_buffers,
+        )
+        # TODO: could be useful to implement a decay warmup via custom avg_fn
 
     def on_after_step(self, trainer: BaseTrainer):
         self.ema_model.update_parameters(trainer.model)
 
     def on_load_state_dict(self, trainer: BaseTrainer, state_dict: dict):
         trainer.logger.info("=> Loading EMA model state ...")
-        set_model_state_dict(self.ema_model, state_dict["ema_model"])
+        incompatible_keys = set_model_state_dict(
+            self.ema_model, state_dict["ema_model"]
+        )
+        # This currently doesn't do anything because strict=True is implicit.
+        log_state_dict_incompatible_keys(
+            trainer.logger,
+            incompatible_keys.missing_keys,
+            incompatible_keys.unexpected_keys,
+        )
 
     def on_state_dict(self, trainer: BaseTrainer, state_dict: dict):
         # Note: sadly, we need to keep the AveragedModel wrapper, to save its n_averaged buffer
@@ -653,9 +687,18 @@ class WandbHook(BaseHook):
                 trainer.logger.debug(f"Dry run log. Would log: {data}")
 
     def on_log_images(self, trainer: BaseTrainer, records: dict, dry_run: bool = False):
+        """Note that the final component of each flattened key becomes the image caption,
+        and the remaining components become the W&B key. For example,
+        ``foo/bar/name1`` and ``foo/bar/name2`` are logged as two images under
+        ``foo/bar``, so W&B displays them together in one panel.
+
+        The list is replaced rather than extended when the same W&B key is logged
+        again at the same step. Callers should therefore collect all images for a
+        panel and pass them in a single call per step.
+        """
         if _dist_rank() == 0:
             wandb_data = {}
-            for k, img in flatten_nested_dict({"vis": records}).items():
+            for k, img in flatten_nested_dict(records).items():
                 file_type = self.image_format(k)
                 wandb_data.setdefault("/".join(k[:-1]), []).append(
                     wandb.Image(
@@ -673,7 +716,7 @@ class WandbHook(BaseHook):
                 trainer.logger.debug(f"Dry run log. Would log: {wandb_data}")
 
     @staticmethod
-    def _ensure_jpeg_compatible(img: PILImage, bg_color: tuple = (255, 255, 255)):
+    def _ensure_jpeg_compatible(img: "PILImage", bg_color: tuple = (255, 255, 255)):
         if img.mode in ("RGB", "L"):
             return img
         elif img.mode in ("RGBA", "LA"):
@@ -702,6 +745,103 @@ class WandbHook(BaseHook):
         return None
 
 
+class TensorBoardHook(BaseHook):
+    """Log metrics and images to TensorBoard (rank 0 only).
+
+    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(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.7.0 → trainloop-0.9.0}/src/trainloop/trainer.py

@@ -1,6 +1,7 @@
 import logging
 import time
 import warnings
+from collections.abc import Callable, Iterable, Iterator
 from contextlib import closing, nullcontext
 from logging import Logger
 from numbers import Number
@@ -8,21 +9,20 @@ from pathlib import Path
 from typing import (
     TYPE_CHECKING,
     Any,
-    Callable,
-    Iterable,
-    Iterator,
     TypeAlias,
     Union,
 )
 
 import torch
-import torch.nn as nn
+from torch import nn
 from torch.distributed.checkpoint.state_dict import (
     StateDictOptions,
     get_state_dict,
     set_state_dict,
 )
 
+from .utils import log_state_dict_incompatible_keys
+
 if TYPE_CHECKING:
     from .hooks import BaseHook
 
@@ -154,13 +154,18 @@ class BaseTrainer:
         self.grad_scaler.load_state_dict(training_state["grad_scaler"])
 
         self.logger.info("=> Loading model and optimizer state ...")
-        set_state_dict(
+        incompatible_keys = set_state_dict(
             self.model,
             self.optimizer,
             model_state_dict=state_dict["model"],
             optim_state_dict=training_state["optimizer"],
             options=self.state_dict_options,
         )
+        log_state_dict_incompatible_keys(
+            self.logger,
+            incompatible_keys.missing_keys,
+            incompatible_keys.unexpected_keys,
+        )
 
         self.logger.info("=> Loading hook states ...")
         for h in self.hooks:
@@ -203,8 +208,8 @@ class BaseTrainer:
 
         reset_step_info()
         self.step_info["data_time"] = []
-        non_finite_grad_retry_count = 0
-        self.step_info["non_finite_grad_retry_count"] = non_finite_grad_retry_count
+        non_finite_grad_retries = 0
+        self.step_info["non_finite_grad_retries"] = non_finite_grad_retries
         i_acc = 0
         while i_acc < self.gradient_accumulation_steps:
             is_accumulating = i_acc < self.gradient_accumulation_steps - 1
@@ -278,15 +283,14 @@ class BaseTrainer:
                         )
                     ):
                         if self.max_non_finite_grad_retries is None or (
-                            non_finite_grad_retry_count
-                            < self.max_non_finite_grad_retries
+                            non_finite_grad_retries < 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
+                            non_finite_grad_retries += 1
+                            self.step_info["non_finite_grad_retries"] = (
+                                non_finite_grad_retries
                             )
                             self.logger.warning(
-                                f"Gradient is non-finite. Retrying step {self.step} (retry {non_finite_grad_retry_count}"
+                                f"Gradient is non-finite. Retrying step {self.step} (retry {non_finite_grad_retries}"
                                 + (
                                     f"/{self.max_non_finite_grad_retries})."
                                     if self.max_non_finite_grad_retries is not None

{trainloop-0.7.0 → trainloop-0.9.0}/src/trainloop/utils.py

@@ -1,7 +1,9 @@
 # modified from: https://github.com/microsoft/MoGe/blob/6b8b43db567ca4b08615c39b42cffd6c76cada29/moge/utils/tools.py
 
 import math
-from typing import Any, Generator, MutableMapping
+from collections.abc import Generator, MutableMapping
+from logging import Logger
+from typing import Any
 
 
 def traverse_nested_dict_keys(
@@ -50,7 +52,7 @@ def key_average(list_of_dicts: list, exclude_nan: bool = False) -> dict[str, Any
 
 
 def flatten_nested_dict(
-    d: dict[str, Any], parent_key: tuple[str, ...] = None
+    d: dict[str, Any], parent_key: tuple[str, ...] | None = None
 ) -> dict[tuple[str, ...], Any]:
     """
     Flattens a nested dictionary into a single-level dictionary, with keys as tuples.
@@ -65,3 +67,20 @@ def flatten_nested_dict(
         else:
             items.append((new_key, v))
     return dict(items)
+
+
+def log_state_dict_incompatible_keys(
+    logger: Logger, missing_keys: list[str], unexpected_keys: list[str]
+) -> None:
+    if missing_keys:
+        logger.warning(
+            "Missing keys in state_dict: {}.".format(
+                ", ".join(f'"{key}"' for key in missing_keys)
+            )
+        )
+    if unexpected_keys:
+        logger.warning(
+            "Unexpected keys in state_dict: {}.".format(
+                ", ".join(f'"{key}"' for key in unexpected_keys)
+            )
+        )