torch-tk 1.2

Streamlines training, checkpoint management, and diagnostics of PyTorch models.

Torch ToolKit (torch-tk)

torch-tk streamlines training, checkpoint management, and diagnostics of PyTorch models.

Overview

torch-tk adds a small amount of structure around PyTorch models and optimizers to simplify training and automate saving and restoring checkpoints.

torch-tk provides a model base class, optimizers, a checkpoint manager, a trainer, and diagnostics utilities. The model class inherits from torch.nn.Module, and the torch-tk optimizers are wrappers around torch.optim optimizers such as torch.optim.SGD and torch.optim.Adam. The torch-tk model and optimizer classes thus preserve functionality and interface of PyTorch modules and optimizers.

Key features

torch-tk models and optimizers are self-describing: A model derived from torch_tk.models.Model and the optimizers in torch_tk.optimizers provide all information needed to save their state and recreate the same model and optimizer instances later. In practice, this means a model and optimizer can save to file the constructor arguments and state parameters needed to rebuild them, allowing to load them back into fresh instances.

torch-tk provides a CheckPointManager. The CheckPointManager manages saving, loading, and reconstruing both the model and the optimizer in the state that created the checkpoint. All that is required is that the class paths are available to import the original model and optimizer classes.

torch-tk provides trainer classes for epoch-based PyTorch training. It includes:

• a standard Trainer for targets without missing values,
• a MaskedTrainer for targets that may contain missing values represented by NaN.

Both trainers support training either from a DataLoader or directly from tensors, and record basic diagnostics such as loss and epoch wallclock time.

torch-tk provides a Diagnostics class for storing sample-resolved loss information together with training metadata. These diagnostics can be created from tensors or data loaders and can be saved to and restored from netCDF files for later analysis.

Workflow

The workflow is shown in the torch-tk HowTo Jupyter notebook.

Installation

pip

pip install torch-tk

conda / mamba

mamba install -c jan.kazil -c conda-forge torch-tk

Public API

Classes

• Model

  • A base class which makes models self-describing and automatically reconstructible by the CheckPointManager
  • Automatically rebuilds a model from a saved file

• SGD, Adam, ...

  Wrapper classes for PyTorch optimizers that make the optimizers self-describing and automatically reconstructible by the CheckPointManager

• Trainer

  • Trains a model from
    • a torch.utils.data.DataLoader
    • or directly from tensors, using an efficient batching mechanism
  • Records training loss and model timing per epoch

• MaskedTrainer

  • Trains a model when target tensors may contain missing values represented by NaN
  • Requires a masked loss function that
    • excludes NaN target values before computing the loss
    • returns both the scalar loss and the number of valid prediction-target pairs used in the loss calculation
  • Trains from
    • a torch.utils.data.DataLoader
    • or directly from tensors, using an efficient batching mechanism
  • Skips batches with no valid target values
  • Records training loss, optional validation loss, and model timing per diagnostic epoch

• CheckPointManager

  • Saves and restores model training states
  • Automatically rebuilds both a model and its optimizer from a saved checkpoint file

• Diagnostics

  • Computes, stores, and plots per-sample loss and per-sample loss probability distribution
  • Saves and restores diagnostics in netCDF file format
  • Identifies worst-loss samples

Modules

torch_tk.models.model

Provides the abstract Model base class for models that can describe, save, restore, and reconstruct themselves. The Model class inherits from torch.nn.Module, and thus provides the standard PyTorch Module interface.

The Model class defines and provides the following methods:

• Model.forward(xb): Abstract method that computes the forward pass.
• Model.constructor_dict(): Abstract method that returns the constructor arguments needed to reconstruct the model.
• Model.save_state_dict_to_file(path): Save only the state dictionary.
• Model.save_to_file(path): Save constructor arguments and state dictionary needed to recreate the model.
• Model.load_from_file(path, device=None): Recreate a model from a saved file.
• Model.clone(constructor_dict, state_dict, device=None): Reconstruct a model from constructor arguments and state.

torch_tk.optimizers.sgd

Wrapper around torch.optim.SGD to make it self-describing and automatically reconstructible.

• SGD(...): Subclass of torch.optim.SGD that stores its constructor arguments on the instance.
• SGD.constructor_dict(): Return the stored optimizer constructor settings excluding params.

torch_tk.optimizers.adam

Wrapper around torch.optim.Adam to make it self-describing and automatically reconstructible.

• Adam(...): Subclass of torch.optim.Adam that stores its constructor arguments on the instance.
• Adam.constructor_dict(): Return the stored optimizer constructor settings excluding params.

torch_tk.training.trainer

Provides trainer classes for epoch-based PyTorch training and simple diagnostics.

• Trainer

  • Standard trainer for targets without missing values.

  • Constructor

    • Trainer(model, optimizer, loss_function, epoch=0): Initialize trainer state with a scalar loss function.

  • Training methods

    • Trainer.train_with_dataloader(data_loader, num_epochs, epoch_diag_step=1, valid_data_loader=None, verbose=True): Train from a DataLoader.
    • Trainer.train_with_data(x_train, y_train, bs, num_epochs, epoch_diag_step=1, x_valid=None, y_valid=None, shuffle=True, verbose=True): Train from in-memory tensors.

  • Plotting methods

    • Trainer.plot_loss(...): Plot recorded diagnostic loss versus epoch.
    • Trainer.plot_wallclock_time(...): Plot recorded epoch wallclock time versus epoch.

• MaskedTrainer

  • Trainer for targets that may contain missing values represented by NaN.

  • Constructor

    • MaskedTrainer(model, optimizer, masked_loss_function, epoch=0): Initialize trainer state with a masked loss function.

  • Required masked-loss behavior

    • Exclude NaN target values before subtraction, squaring, or other loss operations.
    • Return (loss, target_valid_n), where target_valid_n is the number of valid prediction-target pairs used in the loss calculation.
    • Allow batches with target_valid_n == 0 to be skipped.

  • Training methods

    • MaskedTrainer.train_with_dataloader(data_loader, num_epochs, epoch_diag_step=1, valid_data_loader=None, verbose=True): Train from a DataLoader, skipping batches with no valid target values.
    • MaskedTrainer.train_with_data(x_train, y_train, bs, num_epochs, epoch_diag_step=1, x_valid=None, y_valid=None, shuffle=True, verbose=True): Train from in-memory tensors, skipping batches with no valid target values.

  • Plotting methods

    • MaskedTrainer.plot_loss(...): Plot recorded diagnostic training loss and, when available, validation loss versus epoch.
    • MaskedTrainer.plot_wallclock_time(...): Plot recorded epoch wallclock time versus epoch.

torch_tk.checkpoints.checkpoint_manager

Provides checkpoint management for saving and reconstructing a model and optimizer together.

• CheckPointManager(model, optimizer, directory): Manage checkpoint saving in a directory.
• CheckPointManager.save(epoch, batch_size): Save a checkpoint containing epoch, class paths, constructor dictionaries, and state dictionaries.
• CheckPointManager.load_from_file(file_path, device=None): Reconstruct and return checkpoint_manager, model, optimizer, epoch, batch_size from a checkpoint file.

torch_tk.diagnostics.loss

Provides utilities for computing per-sample loss.

• per_sample_loss_from_data_loader(model, loss_function_sample_resolved, data_loader): Compute per-sample losses and their mean from a DataLoader.
• per_sample_loss_from_data(model, loss_function_sample_resolved, x_data, y_data, chunk_size=None): Compute per-sample losses and their mean from in-memory tensors.
• model_worst_loss(model, loss_function_sample_resolved, x_data, y_data, n, chunk_size=None): Return the indices and values of the n worst losses.

torch_tk.diagnostics.diagnostics

Provides the Diagnostics class for sample-resolved loss diagnostics and analysis.

• Diagnostics.from_data_loader(...): Build a diagnostics object from a model evaluated on a DataLoader.
• Diagnostics.from_data(...): Build a diagnostics object from in-memory tensors.
• Diagnostics.from_netcdf(path): Restore a diagnostics object from a saved netCDF file.
• Diagnostics(...): Construct a diagnostics object from metadata, epochs, and per-sample loss data.
• Diagnostics.__add__(other): Concatenate compatible diagnostics across epochs.
• Diagnostics.to_netcdf(directory, verbose=True): Save diagnostics to a netCDF file.

torch_tk.diagnostics.plotting

Provides utilities for plotting diagnostics.

• plot_positive_loss_kde_pdf(diagnostics, plot_file=None, title=None, font_factor=1.5, figsize=(9, 6), xlim=None, ylim=None, loss_name='Loss', xlog=False, ylog=False, bin_n=100, per_dlog10=False, show_plot=True, verbose=True, epoch_skip=1): Plot kernel-density-estimated probability distribution functions (PDFs) of positive loss values across one or more diagnostics objects and epochs. Optionally plot density per unit log10(loss) instead of per unit loss. All loss values must be positive.

• plot_positive_loss_hist_pdf(diagnostics, plot_file=None, title=None, font_factor=1.5, figsize=(9, 6), xlim=None, ylim=None, loss_name='Loss', xlog=False, ylog=False, bin_n=25, per_dlog10=False, show_plot=True, verbose=True, epoch_skip=1): Plot histogram-based probability distribution functions (PDFs) of positive loss values across one or more diagnostics objects and epochs using logarithmically spaced bins. Optionally plot density per unit log10(loss) instead of per unit loss. All loss values must be positive.

• plot_positive_loss_hist_1st_moment_density(diagnostics, plot_file=None, title=None, font_factor=1.5, figsize=(9, 6), xlim=None, ylim=None, loss_name='Loss', xlog=False, ylog=False, bin_n=25, per_dlog10=False, show_plot=True, verbose=True, epoch_skip=1): Plot histogram-based first-moment densities of positive loss values across one or more diagnostics objects and epochs using logarithmically spaced bins. By default, this plots loss · dP/d(loss); optionally it can plot loss · dP/dlog10(loss) for logarithmic comparison across scales. All loss values must be positive.

Notes and limitations

• The checkpoint mechanism assumes that models and optimizers are importable from stable class paths and expose constructor_dict(), state_dict(), and load_state_dict().
• The checkpoint design is not suitable for optimizers that require non-serializable constructor inputs or custom parameter-group reconstruction beyond model.parameters().
• The recorded epoch loss in Trainer is exact only when the supplied loss function returns the mean per-sample loss over each batch, as stated in the trainer docstrings.

Development

Code Quality and Testing Commands

• make fmt - Runs ruff format, which reformats Python files according to the style rules in pyproject.toml.
• make lint - Runs ruff check --fix, which lints the code and auto-fixes what it can.
• make check - Runs formatting and linting.
• make type - Currently disabled. Intended to run mypy using the settings in pyproject.toml.
• make test - Runs pytest with the test settings configured in pyproject.toml.

Author

Jan Kazil - [e-mail] - jankazil.com

License

BSD-3-Clause