xax 0.0.3__py3-none-any.whl → 0.0.5__py3-none-any.whl

xax/task/mixins/logger.py

@@ -4,14 +4,12 @@ import os
 from dataclasses import dataclass
 from pathlib import Path
 from types import TracebackType
-from typing import Callable, Generic, Self, Sequence, TypeVar
-
-from jaxtyping import Array
+from typing import Generic, Self, TypeVar
 
 from xax.core.conf import Device as BaseDeviceConfig, field
 from xax.core.state import State
 from xax.task.base import BaseConfig, BaseTask
-from xax.task.logger import ChannelSelectMode, Logger, LoggerImpl, Number
+from xax.task.logger import Logger, LoggerImpl
 from xax.task.loggers.json import JsonLogger
 from xax.task.loggers.state import StateLogger
 from xax.task.loggers.stdout import StdoutLogger
@@ -59,252 +57,6 @@ class LoggerMixin(BaseTask[Config], Generic[Config]):
     def write_logs(self, state: State) -> None:
         self.logger.write(state)
 
-    def log_scalar(self, key: str, value: Callable[[], Number] | Number, *, namespace: str | None = None) -> None:
-        self.logger.log_scalar(key, value, namespace=namespace)
-
-    def log_string(self, key: str, value: Callable[[], str] | str, *, namespace: str | None = None) -> None:
-        self.logger.log_string(key, value, namespace=namespace)
-
-    def log_image(
-        self,
-        key: str,
-        value: Callable[[], Array] | Array,
-        *,
-        namespace: str | None = None,
-        keep_resolution: bool = False,
-    ) -> None:
-        self.logger.log_image(
-            key,
-            value,
-            namespace=namespace,
-            keep_resolution=keep_resolution,
-        )
-
-    def log_labeled_image(
-        self,
-        key: str,
-        value: Callable[[], tuple[Array, str]] | tuple[Array, str],
-        *,
-        namespace: str | None = None,
-        max_line_length: int | None = None,
-        keep_resolution: bool = False,
-        centered: bool = True,
-    ) -> None:
-        self.logger.log_labeled_image(
-            key,
-            value,
-            namespace=namespace,
-            max_line_length=max_line_length,
-            keep_resolution=keep_resolution,
-            centered=centered,
-        )
-
-    def log_images(
-        self,
-        key: str,
-        value: Callable[[], Array] | Array,
-        *,
-        namespace: str | None = None,
-        keep_resolution: bool = False,
-        max_images: int | None = None,
-        sep: int = 0,
-    ) -> None:
-        self.logger.log_images(
-            key,
-            value,
-            namespace=namespace,
-            keep_resolution=keep_resolution,
-            max_images=max_images,
-            sep=sep,
-        )
-
-    def log_labeled_images(
-        self,
-        key: str,
-        value: Callable[[], tuple[Array, Sequence[str]]] | tuple[Array, Sequence[str]],
-        *,
-        namespace: str | None = None,
-        max_line_length: int | None = None,
-        keep_resolution: bool = False,
-        max_images: int | None = None,
-        sep: int = 0,
-        centered: bool = True,
-    ) -> None:
-        self.logger.log_labeled_images(
-            key,
-            value,
-            namespace=namespace,
-            max_line_length=max_line_length,
-            keep_resolution=keep_resolution,
-            max_images=max_images,
-            sep=sep,
-            centered=centered,
-        )
-
-    def log_audio(
-        self,
-        key: str,
-        value: Callable[[], Array] | Array,
-        *,
-        namespace: str | None = None,
-        sample_rate: int = 44100,
-        log_spec: bool = True,
-        n_fft_ms: float = 32.0,
-        hop_length_ms: float | None = None,
-        channel_select_mode: ChannelSelectMode = "first",
-        keep_resolution: bool = False,
-    ) -> None:
-        self.logger.log_audio(
-            key,
-            value,
-            namespace=namespace,
-            sample_rate=sample_rate,
-            log_spec=log_spec,
-            n_fft_ms=n_fft_ms,
-            hop_length_ms=hop_length_ms,
-            channel_select_mode=channel_select_mode,
-            keep_resolution=keep_resolution,
-        )
-
-    def log_audios(
-        self,
-        key: str,
-        value: Callable[[], Array] | Array,
-        *,
-        namespace: str | None = None,
-        sep_ms: float = 0.0,
-        max_audios: int | None = None,
-        sample_rate: int = 44100,
-        log_spec: bool = True,
-        n_fft_ms: float = 32.0,
-        hop_length_ms: float | None = None,
-        channel_select_mode: ChannelSelectMode = "first",
-        spec_sep: int = 0,
-        keep_resolution: bool = False,
-    ) -> None:
-        self.logger.log_audios(
-            key,
-            value,
-            namespace=namespace,
-            sep_ms=sep_ms,
-            max_audios=max_audios,
-            sample_rate=sample_rate,
-            log_spec=log_spec,
-            n_fft_ms=n_fft_ms,
-            hop_length_ms=hop_length_ms,
-            channel_select_mode=channel_select_mode,
-            spec_sep=spec_sep,
-            keep_resolution=keep_resolution,
-        )
-
-    def log_spectrogram(
-        self,
-        key: str,
-        value: Callable[[], Array] | Array,
-        *,
-        namespace: str | None = None,
-        sample_rate: int = 44100,
-        n_fft_ms: float = 32.0,
-        hop_length_ms: float | None = None,
-        channel_select_mode: ChannelSelectMode = "first",
-        keep_resolution: bool = False,
-    ) -> None:
-        self.logger.log_spectrogram(
-            key,
-            value,
-            namespace=namespace,
-            sample_rate=sample_rate,
-            n_fft_ms=n_fft_ms,
-            hop_length_ms=hop_length_ms,
-            channel_select_mode=channel_select_mode,
-            keep_resolution=keep_resolution,
-        )
-
-    def log_spectrograms(
-        self,
-        key: str,
-        value: Callable[[], Array] | Array,
-        *,
-        namespace: str | None = None,
-        max_audios: int | None = None,
-        sample_rate: int = 44100,
-        n_fft_ms: float = 32.0,
-        hop_length_ms: float | None = None,
-        channel_select_mode: ChannelSelectMode = "first",
-        spec_sep: int = 0,
-        keep_resolution: bool = False,
-    ) -> None:
-        self.logger.log_spectrograms(
-            key,
-            value,
-            namespace=namespace,
-            max_audios=max_audios,
-            sample_rate=sample_rate,
-            n_fft_ms=n_fft_ms,
-            hop_length_ms=hop_length_ms,
-            channel_select_mode=channel_select_mode,
-            spec_sep=spec_sep,
-            keep_resolution=keep_resolution,
-        )
-
-    def log_video(
-        self,
-        key: str,
-        value: Callable[[], Array] | Array,
-        *,
-        namespace: str | None = None,
-        fps: int | None = None,
-        length: float | None = None,
-    ) -> None:
-        self.logger.log_video(
-            key,
-            value,
-            namespace=namespace,
-            fps=fps,
-            length=length,
-        )
-
-    def log_videos(
-        self,
-        key: str,
-        value: Callable[[], Array | list[Array]] | Array | list[Array],
-        *,
-        namespace: str | None = None,
-        max_videos: int | None = None,
-        sep: int = 0,
-        fps: int | None = None,
-        length: int | None = None,
-    ) -> None:
-        self.logger.log_videos(
-            key,
-            value,
-            namespace=namespace,
-            max_videos=max_videos,
-            sep=sep,
-            fps=fps,
-            length=length,
-        )
-
-    def log_histogram(self, key: str, value: Callable[[], Array] | Array, *, namespace: str | None = None) -> None:
-        self.logger.log_histogram(key, value, namespace=namespace)
-
-    def log_point_cloud(
-        self,
-        key: str,
-        value: Callable[[], Array] | Array,
-        *,
-        namespace: str | None = None,
-        max_points: int = 1000,
-        colors: Callable[[], Array] | Array | None = None,
-    ) -> None:
-        self.logger.log_point_cloud(
-            key,
-            value,
-            namespace=namespace,
-            max_points=max_points,
-            colors=colors,
-        )
-
     def __enter__(self) -> Self:
         self.logger.__enter__()
         return self

xax/task/mixins/process.py

@@ -38,6 +38,10 @@ class ProcessMixin(BaseTask[Config], Generic[Config]):
     def multiprocessing_context(self) -> BaseContext:
         return self._mp_ctx
 
+    @property
+    def multiprocessing_manager(self) -> SyncManager:
+        return self._mp_manager
+
     def on_training_end(self, state: State) -> State:
         state = super().on_training_end(state)
 

xax/task/mixins/train.py

@@ -1,9 +1,11 @@
 """Defines a mixin for running the training loop."""
 
+import bdb
 import contextlib
 import functools
 import itertools
 import logging
+import signal
 import sys
 import textwrap
 import time
@@ -11,20 +13,21 @@ import traceback
 from abc import ABC, abstractmethod
 from dataclasses import dataclass, is_dataclass
 from threading import Thread
-from typing import Generic, Literal, Mapping, Sequence, TypeVar, cast, get_args
+from typing import Any, Generic, Literal, Mapping, Sequence, TypeVar, cast, get_args
 
 import equinox as eqx
 import jax
 import jax.numpy as jnp
 import numpy as np
 import optax
-from jaxtyping import Array
+from jaxtyping import Array, PyTree
 from omegaconf import DictConfig
 
 from xax.core.conf import field
 from xax.core.state import Phase, State
 from xax.nn.parallel import is_master
 from xax.task.mixins.artifacts import ArtifactsConfig, ArtifactsMixin
+from xax.task.mixins.checkpointing import CheckpointingConfig, CheckpointingMixin
 from xax.task.mixins.data_loader import DataloadersConfig, DataloadersMixin
 from xax.task.mixins.logger import LoggerConfig, LoggerMixin
 from xax.task.mixins.runnable import RunnableConfig, RunnableMixin
@@ -32,6 +35,8 @@ from xax.task.mixins.step_wrapper import StepContextConfig, StepContextMixin
 from xax.utils.experiments import (
     StateTimer,
     TrainingFinishedError,
+    diff_configs,
+    get_diff_string,
     get_git_state,
     get_training_code,
 )
@@ -40,9 +45,11 @@ from xax.utils.text import highlight_exception_message, show_info
 
 logger = logging.getLogger(__name__)
 
-Model = TypeVar("Model", bound=eqx.Module)
-Batch = TypeVar("Batch")
-Output = TypeVar("Output")
+# Batch = TypeVar("Batch")
+# Output = TypeVar("Output")
+
+Batch = Any
+Output = Any
 
 StepKind = Literal["step", "sample", "second"]
 
@@ -125,6 +132,7 @@ class ValidStepTimer:
 
 @dataclass
 class TrainConfig(
+    CheckpointingConfig,
     DataloadersConfig,
     LoggerConfig,
     StepContextConfig,
@@ -145,12 +153,13 @@ Config = TypeVar("Config", bound=TrainConfig)
 
 
 class TrainMixin(
+    CheckpointingMixin[Config],
     DataloadersMixin[Config],
     LoggerMixin[Config],
     StepContextMixin[Config],
     ArtifactsMixin[Config],
     RunnableMixin[Config],
-    Generic[Config, Model, Batch, Output],
+    Generic[Config],
     ABC,
 ):
     valid_step_timer: ValidStepTimer
@@ -159,7 +168,6 @@ class TrainMixin(
     _training_over_flag: bool
     _last_printed_remaining_time: float
     _step_kind: StepKind
-    _prng_key: jnp.ndarray
 
     def __init__(self, config: Config) -> None:
         super().__init__(config)
@@ -183,12 +191,8 @@ class TrainMixin(
         # The kind of step that was specified in the config.
         self._step_kind = cast_step_kind(self.config.step_kind)
 
-        # Defines a PRNG key for the task.
-        self._prng_key = jax.random.PRNGKey(self.config.random_seed)
-
-    @property
     def prng_key(self) -> jnp.ndarray:
-        return self._prng_key
+        return jax.random.PRNGKey(self.config.random_seed)
 
     def on_step_end(self, state: State) -> State:
         state = super().on_step_end(state)
@@ -198,7 +202,7 @@ class TrainMixin(
             },
         )
 
-    def log_train_step(self, model: Model, batch: Batch, output: Output, state: State) -> None:
+    def log_train_step(self, model: PyTree, batch: Batch, output: Output, state: State) -> None:
         """Override this function to do logging during the training phase.
 
         This function is called after the model forward pass and before the
@@ -211,7 +215,7 @@ class TrainMixin(
             state: The current training state.
         """
 
-    def log_valid_step(self, model: Model, batch: Batch, output: Output, state: State) -> None:
+    def log_valid_step(self, model: PyTree, batch: Batch, output: Output, state: State) -> None:
         """Override this function to do logging during the validation phase.
 
         This function is called after the model forward pass. It is called in
@@ -224,7 +228,7 @@ class TrainMixin(
             state: The current training state.
         """
 
-    def log_step(self, model: Model, batch: Batch, output: Output, state: State) -> None:
+    def log_step(self, model: PyTree, batch: Batch, output: Output, state: State) -> None:
         phase = state.phase
 
         # Log the state timers.
@@ -232,7 +236,7 @@ class TrainMixin(
         timer.step(state)
         for ns, d in timer.log_dict().items():
             for k, v in d.items():
-                self.log_scalar(k, v, namespace=ns)
+                self.logger.log_scalar(k, v, namespace=ns)
 
         # Delegate to the appropriate logging function based on the phase.
         match phase:
@@ -244,7 +248,7 @@ class TrainMixin(
                 raise KeyError(f"Unknown phase: {phase}")
 
     @abstractmethod
-    def get_model(self) -> Model:
+    def get_model(self) -> PyTree:
         """Returns the Equinox model to train.
 
         Returns:
@@ -259,11 +263,34 @@ class TrainMixin(
             The optimizer to use to train the model.
         """
 
-    def get_initial_opt_state(self, model: Model, optimizer: optax.GradientTransformation) -> optax.OptState:
+    def get_initial_opt_state(self, model: PyTree, optimizer: optax.GradientTransformation) -> optax.OptState:
         return optimizer.init(eqx.filter(model, eqx.is_array))
 
+    def load_initial_state(self) -> tuple[PyTree, optax.GradientTransformation, optax.OptState, State]:
+        init_ckpt_path = self.get_init_ckpt_path()
+
+        if init_ckpt_path is not None:
+            logger.info("Loading checkpoint from %s", init_ckpt_path)
+            with self.step_context("load_checkpoint"):
+                model, optimizer, opt_state, state, config = self.load_checkpoint(init_ckpt_path)
+                config_diff = get_diff_string(diff_configs(config, cast(DictConfig, self.config)))
+                if config_diff:
+                    logger.warning("Loaded config differs from current config:\n%s", config_diff)
+                return model, optimizer, opt_state, state
+
+        with self.step_context("get_model"):
+            model = self.get_model()
+
+        with self.step_context("get_optimizer"):
+            optimizer = self.get_optimizer()
+
+        with self.step_context("get_initial_opt_state"):
+            opt_state = optimizer.init(eqx.filter(model, eqx.is_array))
+
+        return model, optimizer, opt_state, State.init_state()
+
     @abstractmethod
-    def get_output(self, model: Model, batch: Batch, state: State) -> Output:
+    def get_output(self, model: PyTree, batch: Batch, state: State) -> Output:
         """Gets the output from the model.
 
         By default, we assume the model is a function that takes the batch as
@@ -276,7 +303,7 @@ class TrainMixin(
             state: The current training state.
         """
 
-    def compute_loss(self, model: Model, batch: Batch, output: Output, state: State) -> Array:
+    def compute_loss(self, model: PyTree, batch: Batch, output: Output, state: State) -> Array:
         """Gets the loss for the current batch.
 
         By default, we assume the model is a function that takes the batch as
@@ -296,7 +323,7 @@ class TrainMixin(
             raise ValueError(f"When model output is not the loss, you must override `compute_loss`. Got {type(output)}")
         return output
 
-    def get_output_and_loss(self, model: Model, batch: Batch, state: State) -> tuple[Array, Output]:
+    def get_output_and_loss(self, model: PyTree, batch: Batch, state: State) -> tuple[Array, Output]:
         output = self.get_output(model, batch, state)
         loss = self.compute_loss(model, batch, output, state)
         return loss, output
@@ -304,12 +331,12 @@ class TrainMixin(
     @eqx.filter_jit
     def update(
         self,
-        model: Model,
+        model: PyTree,
         optimizer: optax.GradientTransformation,
         opt_state: optax.OptState,
         batch: Batch,
         state: State,
-    ) -> tuple[Array, Model, optax.OptState, Output]:
+    ) -> tuple[Array, PyTree, optax.OptState, Output]:
         (loss, output), grads = eqx.filter_value_and_grad(self.get_output_and_loss, has_aux=True)(model, batch, state)
         updates, opt_state = optimizer.update(grads, opt_state)
         model = eqx.apply_updates(model, updates)
@@ -358,7 +385,7 @@ class TrainMixin(
         remaining_percent = self.get_remaining_percent(state)
         if remaining_percent is None:
             return False
-        self.log_scalar("percent", remaining_percent, namespace="⏰ remaining")
+        self.logger.log_scalar("percent", remaining_percent, namespace="⏰ remaining")
         self.maybe_log_termination_time(remaining_percent, state)
         return remaining_percent <= 0.0
 
@@ -387,15 +414,15 @@ class TrainMixin(
 
     def train_step(
         self,
-        model: Model,
+        model: PyTree,
         optimizer: optax.GradientTransformation,
         opt_state: optax.OptState,
         batch: Batch,
         state: State,
-    ) -> tuple[Model, optax.OptState, State]:
+    ) -> tuple[PyTree, optax.OptState, State]:
         state = state.with_phase("train")
         loss, model, opt_state, output = self.update(model, optimizer, opt_state, batch, state)
-        self.log_scalar("loss", loss, namespace="loss")
+        self.logger.log_scalar("loss", loss, namespace="loss")
         self.log_step(model, batch, output, state)
         self.write_logs(state)
         return (
@@ -409,10 +436,10 @@ class TrainMixin(
             ),
         )
 
-    def val_step(self, model: Model, batch: Batch, state: State) -> tuple[Model, State]:
+    def val_step(self, model: PyTree, batch: Batch, state: State) -> tuple[PyTree, State]:
         state = state.with_phase("valid")
         loss, output = eqx.filter_jit(self.get_output_and_loss)(model, batch, state)
-        self.log_scalar("loss", loss, namespace="loss")
+        self.logger.log_scalar("loss", loss, namespace="loss")
         self.log_step(model, batch, output, state)
         self.write_logs(state)
         return model, state.replace(
@@ -462,20 +489,15 @@ class TrainMixin(
             ctx.enter_context(train_pf)
             ctx.enter_context(valid_pf)
 
-            # Gets the model.
-            with self.step_context("get_model"):
-                model = self.get_model()
+            model, optimizer, opt_state, state = self.load_initial_state()
 
-            # Gets the optimizer.
-            with self.step_context("get_optimizer"):
-                optimizer = self.get_optimizer()
+            state = self.on_training_start(state)
 
-            # Gets the initial optimizer state.
-            with self.step_context("get_initial_opt_state"):
-                opt_state = optimizer.init(eqx.filter(model, eqx.is_array))
+            def on_exit() -> None:
+                self.save_checkpoint(model, optimizer, opt_state, state)
 
-            state = State.init_state()
-            state = self.on_training_start(state)
+            # Handle user-defined interrupts during the training loop.
+            self.add_signal_handler(on_exit, signal.SIGUSR1, signal.SIGTERM)
 
             try:
                 while True:
@@ -494,17 +516,26 @@ class TrainMixin(
                         with self.step_context("on_step_end"):
                             state = self.on_step_end(state)
 
+                        if self.should_checkpoint(state):
+                            self.save_checkpoint(model, optimizer, opt_state, state)
+
             except TrainingFinishedError:
                 if is_master():
                     show_info(
                         f"Finished training after {state.num_steps} steps, {state.num_samples} samples",
                         important=True,
                     )
+                self.save_checkpoint(model, optimizer, opt_state, state)
+
+            except (KeyboardInterrupt, bdb.BdbQuit):
+                if is_master():
+                    show_info("Interrupted training", important=True)
 
             except BaseException:
                 exception_tb = textwrap.indent(highlight_exception_message(traceback.format_exc()), "  ")
                 sys.stdout.write(f"Caught exception during training loop:\n\n{exception_tb}\n")
                 sys.stdout.flush()
+                self.save_checkpoint(model, optimizer, opt_state, state)
 
             finally:
                 state = self.on_training_end(state)

xax/task/task.py

@@ -7,7 +7,8 @@ from xax.task.base import BaseConfig, BaseTask
 from xax.task.mixins import (
     ArtifactsConfig,
     ArtifactsMixin,
-    Batch,
+    CheckpointingConfig,
+    CheckpointingMixin,
     CPUStatsConfig,
     CPUStatsMixin,
     DataloadersConfig,
@@ -16,8 +17,6 @@ from xax.task.mixins import (
     GPUStatsMixin,
     LoggerConfig,
     LoggerMixin,
-    Model,
-    Output,
     ProcessConfig,
     ProcessMixin,
     RunnableConfig,
@@ -32,6 +31,7 @@ from xax.task.mixins import (
 @dataclass
 class Config(
     TrainConfig,
+    CheckpointingConfig,
     DataloadersConfig,
     CPUStatsConfig,
     GPUStatsConfig,
@@ -49,7 +49,8 @@ ConfigT = TypeVar("ConfigT", bound=Config)
 
 
 class Task(
-    TrainMixin[ConfigT, Model, Batch, Output],
+    TrainMixin[ConfigT],
+    CheckpointingMixin[ConfigT],
     DataloadersMixin[ConfigT],
     CPUStatsMixin[ConfigT],
     GPUStatsMixin[ConfigT],
@@ -59,6 +60,6 @@ class Task(
     ArtifactsMixin[ConfigT],
     RunnableMixin[ConfigT],
     BaseTask[ConfigT],
-    Generic[ConfigT, Model, Batch, Output],
+    Generic[ConfigT],
 ):
     pass

xax/utils/data/collate.py

@@ -167,9 +167,9 @@ def collate(
     # Collate dictionaries if they have the same keys.
     if isinstance(item, dict) and all(set(i.keys()) == set(item.keys()) for i in items):
         output_dict = {}
-        item_keys = set(item.keys())
-        for key in item_keys:
-            output_dict[key] = collate([i[key] for i in items], mode=mode, pad=pad)
+        item_keys_set = set(item.keys())
+        for key_in_set in item_keys_set:
+            output_dict[key_in_set] = collate([i[key_in_set] for i in items], mode=mode, pad=pad)
         return output_dict
 
     # Collate lists and tuples if they have the same lengths.
@@ -186,9 +186,9 @@ def collate(
     # Handles dataclasses.
     if is_dataclass(item):
         output_dict = {}
-        item_keys = item.__dict__.keys()
-        for key in item_keys:
-            output_dict[key] = collate([getattr(i, key) for i in items], mode=mode, pad=pad)
+        item_keys_dict = item.__dict__.keys()
+        for key_in_dict in item_keys_dict:
+            output_dict[key_in_dict] = collate([getattr(i, key_in_dict) for i in items], mode=mode, pad=pad)
         return item.__class__(**output_dict)
 
     # By default, don't do anything.