torch-tk 1.0.8__py3-none-any.whl

torch_tk/__init__.py

@@ -0,0 +1,29 @@
+from importlib.metadata import PackageNotFoundError, version
+
+_DIST_NAME = "torch-tk"
+
+try:
+    __version__ = version(_DIST_NAME)
+except PackageNotFoundError:
+    pkg = __package__ or __name__.split(".", 1)[0]
+    try:
+        __version__ = version(pkg)
+    except PackageNotFoundError:
+        __version__ = "0.0.0+local"
+
+from .models.model import Model
+from .optimizers.sgd import SGD
+from .optimizers.adam import Adam
+from .training.trainer import Trainer
+from .checkpoints.checkpoint_manager import CheckPointManager
+from .diagnostics.diagnostics import Diagnostics
+
+__all__ = [
+    "__version__",
+    "Model",
+    "SGD",
+    "Adam",
+    "Trainer",
+    "CheckPointManager",
+    "Diagnostics",
+]

torch_tk/checkpoints/__init__.py

@@ -0,0 +1 @@
+from .checkpoint_manager import CheckPointManager

torch_tk/checkpoints/checkpoint_manager.py

@@ -0,0 +1,139 @@
+'''
+Checkpoint utilities for saving and restoring self-describing model and optimizer state.
+
+This module provides the CheckPointManager class, which saves checkpoints that
+contain enough information to reconstruct both a model and its optimizer in the
+state that created the checkpoint. Each checkpoint stores the fully qualified
+class path, constructor arguments from constructor_dict(), and state from
+state_dict() for both objects.
+
+Models and optimizers must be importable from a stable class path and must expose
+
+- constructor_dict()
+- state_dict()
+- load_state_dict()
+
+A model must be reconstructible from its class path and constructor_dict().
+An optimizer must be reconstructible from its class path, model.parameters(),
+and constructor_dict().
+
+The constructor data must be serializable. In practice, this means it should
+contain only standard serializable Python values such as numbers, strings,
+lists, tuples, and dictionaries.
+
+This design is not suitable for optimizers that depend on non-serializable
+constructor inputs, non-standard constructor signatures, custom parameter-group
+reconstruction beyond model.parameters(), or runtime state not captured by
+state_dict().
+'''
+
+from pathlib import Path
+
+import torch
+
+from .utils import class_path_of_instance, import_class
+
+
+class CheckPointManager:
+    '''
+    Save checkpoints and rebuild a model and optimizer from a checkpoint file.
+
+    The checkpoint contains the epoch, class paths, constructor arguments, and
+    state dictionaries for the model and optimizer.
+    '''
+
+    def __init__(self, model, optimizer, directory):
+        '''
+        Store the model, optimizer, and checkpoint directory.
+
+        The directory is converted to a Path and created if needed.
+        '''
+        if not isinstance(directory, Path):
+            directory = Path(directory)
+
+        self.model = model
+        self.optimizer = optimizer
+        self.directory = directory
+
+    def save(self, epoch):
+        '''
+        Save a checkpoint for the current model and optimizer state.
+
+        Returns the path to the written checkpoint file.
+        '''
+        if not hasattr(self.model, 'constructor_dict'):
+            raise TypeError('Model must implement constructor_dict().')
+
+        if not hasattr(self.optimizer, 'constructor_dict'):
+            raise TypeError('Optimizer must implement constructor_dict().')
+
+        checkpoint = {
+            'epoch': epoch,
+            'model_class_path': class_path_of_instance(self.model),
+            'model_constructor_dict': self.model.constructor_dict(),
+            'model_state_dict': self.model.state_dict(),
+            'optimizer_class_path': class_path_of_instance(self.optimizer),
+            'optimizer_constructor_dict': self.optimizer.constructor_dict(),
+            'optimizer_state_dict': self.optimizer.state_dict(),
+        }
+
+        file_name = Path(type(self.model).__name__ + '.' + type(self.optimizer).__name__ + '.epoch=' + str(epoch) + '.pt')
+
+        self.directory.mkdir(parents=True, exist_ok=True)
+
+        file_path = self.directory / file_name
+        torch.save(checkpoint, file_path)
+
+        return file_path
+
+    @classmethod
+    def load_from_file(cls, file_path, device=None):
+        '''
+        Load a checkpoint file and reconstruct the model, optimizer, and manager.
+
+        Returns:
+            checkpoint_manager, model, optimizer, epoch
+
+        If device is given, the checkpoint is loaded onto that device and a
+        saved model constructor argument named 'device' is overridden.
+        '''
+        if not isinstance(file_path, Path):
+            file_path = Path(file_path)
+
+        checkpoint = torch.load(file_path, map_location=device)
+
+        # Reconstruct model
+
+        model_class = import_class(checkpoint['model_class_path'])
+
+        model_constructor_dict = checkpoint['model_constructor_dict']
+        model_args = model_constructor_dict.get('args', [])
+        model_kwargs = dict(model_constructor_dict.get('kwargs', {}))
+
+        if 'device' in model_kwargs and device is not None:
+            model_kwargs['device'] = device
+
+        model = model_class(*model_args, **model_kwargs)
+        model.load_state_dict(checkpoint['model_state_dict'])
+
+        if device is not None:
+            model = model.to(device)
+
+        # Reconstruct optimizer
+
+        optimizer_class = import_class(checkpoint['optimizer_class_path'])
+
+        optimizer_constructor_dict = checkpoint['optimizer_constructor_dict']
+        optimizer_args = optimizer_constructor_dict.get('args', [])
+        optimizer_kwargs = dict(optimizer_constructor_dict.get('kwargs', {}))
+
+        optimizer = optimizer_class(model.parameters(), *optimizer_args, **optimizer_kwargs)
+
+        optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
+
+        # Create a checkpoint manager instance
+        checkpoint_manager = cls(model, optimizer, file_path.parent)
+
+        epoch = checkpoint['epoch']
+
+        return checkpoint_manager, model, optimizer, epoch

torch_tk/checkpoints/utils.py

@@ -0,0 +1,33 @@
+import importlib
+
+
+def class_path_of_instance(instance):
+    '''
+    Return the fully qualified class path for an instance's class.
+    '''
+    cls = type(instance)
+    return cls.__module__ + '.' + cls.__qualname__
+
+
+def import_class(class_path):
+    '''
+    Import and return a class from a fully qualified class path.
+
+    Nested classes referenced through __qualname__ are supported.
+    '''
+    parts = class_path.split('.')
+
+    for i in range(len(parts), 0, -1):
+        module_name = '.'.join(parts[:i])
+        try:
+            obj = importlib.import_module(module_name)
+            break
+        except ModuleNotFoundError:
+            continue
+    else:
+        raise ImportError(f'Could not import any module prefix from {class_path!r}')
+
+    for attr in parts[i:]:
+        obj = getattr(obj, attr)
+
+    return obj

torch_tk/diagnostics/__init__.py

@@ -0,0 +1,3 @@
+from .diagnostics import Diagnostics
+from .loss import model_worst_loss, per_sample_loss_from_data, per_sample_loss_from_data_loader
+from .plotting import plot_diagnostics

torch_tk/diagnostics/diagnostics.py

@@ -0,0 +1,281 @@
+'''
+Utilities for storing, combining, and serializing sample-resolved training
+diagnostics.
+
+This module defines the Diagnostics class, which can be created from in-memory
+data, a data loader, or a saved netCDF file. It stores per-sample loss values
+together with model and optimizer metadata and can write the diagnostics back
+to netCDF.
+'''
+
+from pathlib import Path
+
+import numpy as np
+import torch
+import xarray as xr
+
+from .loss import per_sample_loss_from_data, per_sample_loss_from_data_loader
+
+
+class Diagnostics:
+    '''
+    Store sample-resolved loss diagnostics together with training metadata.
+
+    An instance holds model and optimizer identifiers, learning rate, batch
+    size, an optional description, one or more epoch values, and the
+    corresponding per-sample loss tensors. It also supports loading from and
+    saving to netCDF.
+    '''
+
+    @classmethod
+    def from_data_loader(
+        cls,
+        model,
+        loss_function_sample_resolved,
+        optimizer,
+        learning_rate,
+        batch_size,
+        data_loader,
+        description: str = None,
+        epoch=0,
+    ):
+        '''
+        Create a Diagnostics instance from a model evaluated on a data loader.
+
+        The method computes per-sample losses for all samples in the loader and
+        stores them together with model, optimizer, and training metadata.
+        '''
+
+        mean_per_sample_loss, per_sample_loss = per_sample_loss_from_data_loader(
+            model, loss_function_sample_resolved, data_loader
+        )
+
+        return cls(
+            type(model).__name__,
+            type(optimizer).__name__,
+            learning_rate,
+            batch_size,
+            epoch=epoch,
+            per_sample_loss=per_sample_loss,
+            description=description,
+        )
+
+    @classmethod
+    def from_data(
+        cls,
+        model,
+        loss_function_sample_resolved,
+        optimizer,
+        learning_rate,
+        batch_size,
+        x_data,
+        y_data,
+        description: str = None,
+        epoch=0,
+        chunk_size=None,
+    ):
+        '''
+        Create a Diagnostics instance from in-memory input and target tensors.
+
+        The method computes per-sample losses from the supplied tensors,
+        optionally in chunks to reduce memory usage, and stores them
+        together with model, optimizer, and training metadata.
+        '''
+
+        mean_per_sample_loss, per_sample_loss = per_sample_loss_from_data(
+            model, loss_function_sample_resolved, x_data, y_data, chunk_size=chunk_size
+        )
+
+        return cls(
+            type(model).__name__,
+            type(optimizer).__name__,
+            learning_rate,
+            batch_size,
+            epoch=epoch,
+            per_sample_loss=per_sample_loss,
+            description=description,
+        )
+
+    @classmethod
+    def from_netcdf(cls, path):
+        '''
+        Create a Diagnostics instance from a saved netCDF file.
+
+        This restores metadata, epoch values, and per-sample loss data from a
+        previously saved diagnostics file, which is useful when resuming work
+        from a checkpoint.
+        '''
+        if not isinstance(path, Path):
+            path = Path(path)
+
+        ds = xr.open_dataset(path)
+
+        instance = cls(
+            ds.attrs['model'],
+            ds.attrs['optimizer'],
+            ds.attrs['learning_rate'],
+            ds.attrs['batch_size'],
+            description=ds.attrs['description'],
+            per_sample_loss=torch.as_tensor(ds['per_sample_loss'].values),
+            epoch=ds['epoch'].values,
+        )
+
+        ds.close()
+
+        return instance
+
+    def __init__(
+        self, model_name, optimizer_name, learning_rate, batch_size, epoch=None, per_sample_loss=None, description: str = None
+    ):
+        '''
+        Initialize a Diagnostics instance.
+
+        Epoch values and per-sample losses must either both be provided or both
+        be omitted. If per-sample loss is one-dimensional, it is promoted to
+        two dimensions so that the leading dimension corresponds to epoch.
+        '''
+
+        if epoch is None and per_sample_loss is not None:
+            raise ValueError('epoch and per_sample_loss must both be None or both not None.')
+        if epoch is not None and per_sample_loss is None:
+            raise ValueError('epoch and per_sample_loss must both be None or both not None.')
+
+        self.model = model_name
+        self.optimizer = optimizer_name
+
+        self.learning_rate = str(learning_rate)
+        self.batch_size = batch_size
+        self.description = description
+
+        if epoch is not None:
+            self.epoch = np.atleast_1d(np.asarray(epoch, dtype=np.int64))
+        else:
+            self.epoch = None
+
+        # Add a first dimension that corresponds to the value(s) in the list self.epoch
+        if per_sample_loss is not None:
+            if per_sample_loss.ndim == 1:
+                self.per_sample_loss = per_sample_loss.unsqueeze(0)
+            else:
+                self.per_sample_loss = per_sample_loss
+        else:
+            self.per_sample_loss = None
+
+        # Check for consistency
+        if epoch is not None and len(self.epoch) != self.per_sample_loss.shape[0]:
+            raise ValueError('Number of epochs and number of per-sample loss instances do not match.')
+
+        return
+
+    def __add__(self, other):
+        if not isinstance(other, Diagnostics):
+            return NotImplemented
+
+        '''
+        Combine two compatible diagnostics objects.
+
+        The two objects must have matching metadata. When both contain per-sample
+        loss data, their epoch arrays and per-sample loss tensors are concatenated
+        along the epoch dimension.
+        '''
+
+        result = Diagnostics.__new__(Diagnostics)
+
+        assert self.model == other.model, (
+            f"Cannot add Diagnostics objects with different models: {self.model!r} != {other.model!r}"
+        )
+        assert self.optimizer == other.optimizer, (
+            f"Cannot add Diagnostics objects with different optimizers: {self.optimizer!r} != {other.optimizer!r}"
+        )
+        assert self.learning_rate == other.learning_rate, (
+            f"Cannot add Diagnostics objects with different learning rates: {self.learning_rate!r} != {other.learning_rate!r}"
+        )
+        assert self.batch_size == other.batch_size, (
+            f"Cannot add Diagnostics objects with different batch sizes: {self.batch_size!r} != {other.batch_size!r}"
+        )
+        assert self.description == other.description, (
+            f"Cannot add Diagnostics objects with different descriptions: {self.description!r} != {other.description!r}"
+        )
+
+        if self.per_sample_loss is None:
+            result.model = other.model
+            result.optimizer = other.optimizer
+            result.epoch = other.epoch
+            result.learning_rate = other.learning_rate
+            result.batch_size = other.batch_size
+            result.description = other.description
+            result.per_sample_loss = other.per_sample_loss
+        elif other.per_sample_loss is None:
+            result.model = self.model
+            result.optimizer = self.optimizer
+            result.epoch = self.epoch
+            result.learning_rate = self.learning_rate
+            result.batch_size = self.batch_size
+            result.description = self.description
+            result.per_sample_loss = self.per_sample_loss
+        else:
+            result.model = self.model
+            result.optimizer = self.optimizer
+            result.epoch = np.concatenate((self.epoch, other.epoch))
+            result.learning_rate = self.learning_rate
+            result.batch_size = self.batch_size
+            result.description = self.description
+            result.per_sample_loss = torch.cat((self.per_sample_loss, other.per_sample_loss), dim=0)
+
+        return result
+
+    def to_netcdf(self, directory, verbose=True):
+        '''
+        Save the diagnostics object to a netCDF file.
+
+        The output file is written under the given directory, which is created
+        if necessary. The saved dataset includes epoch and sample coordinates,
+        metadata attributes, and the per-sample loss array.
+        '''
+
+        if not isinstance(directory, Path):
+            directory = Path(directory)
+
+        file_name = Path(
+            self.model
+            + '.'
+            + self.optimizer
+            + '.'
+            + self.description
+            + '.epoch='
+            + str(self.epoch[0])
+            + '_to_'
+            + str(self.epoch[-1])
+            + '.nc'
+        )
+
+        # Create Xarray dataset
+        ds = xr.Dataset()
+
+        # Coordinates
+        ds.coords['epoch'] = self.epoch
+        ds.coords['sample'] = np.arange(self.per_sample_loss.shape[1])
+
+        # Global attributes
+        ds.attrs['model'] = self.model
+        ds.attrs['optimizer'] = self.optimizer
+        ds.attrs['description'] = self.description
+        ds.attrs['learning_rate'] = self.learning_rate
+        ds.attrs['batch_size'] = self.batch_size
+
+        # 2D variables
+        ds['per_sample_loss'] = (['epoch', 'sample'], self.per_sample_loss.detach().cpu().numpy())
+        ds['per_sample_loss'].attrs['long_name'] = 'Mean per-sample loss'
+
+        # Save
+        file_path = directory / file_name
+        file_path.parent.mkdir(parents=True, exist_ok=True)
+        ds.to_netcdf(file_path, unlimited_dims='epoch')
+
+        # Close
+        ds.close()
+
+        if verbose:
+            print('Saved diagnostics in ', file_path)
+
+        return file_path

torch_tk/diagnostics/loss.py

@@ -0,0 +1,160 @@
+'''
+Loss-related utilities for evaluating models at per-sample resolution.
+
+This module provides functions to compute per-sample losses from either
+a data loader or in-memory tensors, preserve the model's training state
+during evaluation, and identify the samples with the largest losses.
+'''
+
+import torch
+
+from torch_tk.models.utils import get_model_device
+
+
+def per_sample_loss_from_data_loader(model, loss_function_sample_resolved, data_loader):
+    '''
+    Compute per-sample loss values and their mean from a data loader.
+
+    The model is evaluated without gradient tracking, its original training
+    state is restored afterward, and the returned per-sample loss tensor is
+    moved to CPU memory.
+
+    Returns
+    -------
+    float
+        Mean loss over all samples in the data loader.
+    torch.Tensor
+        A 1-D tensor containing the model loss for each sample in the given data loader,
+        always in CPU memory.
+    '''
+
+    device = get_model_device(model)
+
+    was_training = model.training
+
+    model.eval()
+
+    data_n = len(data_loader.dataset)
+
+    per_sample_loss = torch.empty(data_n, device=device)
+
+    with torch.no_grad():
+        ii = 0
+
+        for x, y in data_loader:
+            batch_n = len(x)
+            loss = loss_function_sample_resolved(model(x.to(device)), y.to(device))
+            if loss.ndim != 1:
+                raise ValueError('Per-sample loss function does not produce a 1-dimensional tensor.')
+            if loss.shape[0] != batch_n:
+                raise ValueError('Per-sample loss function produces fewer or more loss values than samples.')
+            per_sample_loss[ii : ii + batch_n] = loss
+            ii += batch_n
+
+        mean_per_sample_loss = per_sample_loss.mean().item()
+
+    model.train(was_training)
+
+    return mean_per_sample_loss, per_sample_loss.detach().cpu()
+
+
+def per_sample_loss_from_data(model, loss_function_sample_resolved, x_data, y_data, chunk_size=None):
+    '''
+    Compute per-sample loss values and their mean from in-memory tensors.
+
+    The model is evaluated without gradient tracking, optionally in chunks to
+    limit memory usage, and its original training state is restored afterward.
+    The returned per-sample loss tensor is moved to CPU memory.
+
+    Arguments
+    ----------
+    model : torch.nn.Module
+        Model used for prediction.
+    loss_function_sample_resolved : callable
+        Function taking (predictions, targets) and returning a 1-D tensor
+        of per-sample losses.
+    x_data : torch.Tensor
+        Input data of shape (N, ...).
+    y_data : torch.Tensor
+        Target data of shape (N, ...).
+    chunk_size : int, optional
+        Number of samples to process at once. If not provided, all samples will be processed at once.
+
+    Returns
+    -------
+    float
+        Mean loss over all samples.
+    torch.Tensor
+        A 1-D tensor containing the loss for each sample, always in CPU memory.
+    '''
+
+    device = get_model_device(model)
+
+    was_training = model.training
+
+    model.eval()
+
+    data_n = x_data.shape[0]
+
+    per_sample_loss = torch.empty(data_n, device=device)
+
+    with torch.no_grad():
+        if chunk_size:
+            for i_start in range(0, data_n, chunk_size):
+                i_end = min(i_start + chunk_size, data_n)
+                xb = x_data[i_start:i_end].to(device)
+                yb = y_data[i_start:i_end].to(device)
+                per_sample_loss[i_start:i_end] = loss_function_sample_resolved(model(xb), yb)
+        else:
+            per_sample_loss = loss_function_sample_resolved(model(x_data.to(device)), y_data.to(device))
+
+    mean_per_sample_loss = per_sample_loss.mean().item()
+
+    if per_sample_loss.ndim != 1 or per_sample_loss.shape[0] != data_n:
+        raise ValueError('loss_function_sample_resolved must return a 1-D tensor with one loss value per sample.')
+
+    model.train(was_training)
+
+    return mean_per_sample_loss, per_sample_loss.detach().cpu()
+
+
+def model_worst_loss(model, loss_function_sample_resolved, x_data, y_data, n, chunk_size=None):
+    '''
+    Return the indices and loss values of the n worst-performing samples.
+
+    Calculates the
+    - the n indices in the inputs x_data for which the given model has the
+      largest loss relative to the reference data y_data
+    - the corresponding loss values
+
+    Arguments
+    ----------
+    model : torch.nn.Module
+        The model to evaluate.
+    x_data : torch.Tensor
+        Input data, with batch dimension first.
+    y_data : torch.Tensor
+        Target data, with batch dimension first.
+    chunk_size : int, optional
+        Number of samples to process at once.
+
+    Returns
+    -------
+    - the n indices in the inputs x_data for which the given model has the
+      largest mean square error relative to the reference data y_data
+    - the corresponding mean square errors
+    '''
+
+    # Model mean squared error for each data sample
+    with torch.no_grad():
+        mean_per_sample_loss, per_sample_loss = per_sample_loss_from_data(
+            model, loss_function_sample_resolved, x_data, y_data, chunk_size=chunk_size
+        )
+
+    # Indices such that per_sample_loss[idxs] is sorted in descending order
+    idxs = torch.argsort(per_sample_loss, descending=True)
+
+    # The indices of the n elements in x_data that produce the largest model mean square error
+    idxs_worst = idxs[:n].tolist()
+
+    return idxs_worst, per_sample_loss[idxs_worst]