torch-tk 1.0.8__tar.gz

torch_tk-1.0.8/LICENSE

@@ -0,0 +1,13 @@
+BSD 3-Clause License
+
+Copyright (c) 2026, Jan Kazil
+
+Redistribution and use in source and binary forms, with or without modification, are permitted provided that the following conditions are met:
+
+1. Redistributions of source code must retain the above copyright notice, this list of conditions and the following disclaimer.
+
+2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the following disclaimer in the documentation and/or other materials provided with the distribution.
+
+3. Neither the name of the copyright holder nor the names of its contributors may be used to endorse or promote products derived from this software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

torch_tk-1.0.8/PKG-INFO

@@ -0,0 +1,197 @@
+Metadata-Version: 2.4
+Name: torch-tk
+Version: 1.0.8
+Summary: Streamlines training, checkpoint management, and diagnostics of PyTorch models.
+Author-email: Jan Kazil <jan.kazil.dev@gmail.com>
+License-Expression: BSD-3-Clause
+Project-URL: Homepage, https://github.com/jankazil/torch-tk
+Project-URL: Repository, https://github.com/jankazil/torch-tk
+Project-URL: Issues, https://github.com/jankazil/torch-tk/issues
+Keywords: pytorch,training,checkpointing,diagnostics,deep-learning
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Science/Research
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Topic :: Scientific/Engineering
+Classifier: Topic :: Software Development :: Libraries :: Python Modules
+Requires-Python: >=3.12
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: torch
+Requires-Dist: numpy
+Requires-Dist: xarray
+Requires-Dist: matplotlib
+Requires-Dist: scipy
+Provides-Extra: dev
+Requires-Dist: build>=1.2; extra == "dev"
+Requires-Dist: twine>=5.1; extra == "dev"
+Requires-Dist: pytest>=8; extra == "dev"
+Requires-Dist: pytest-cov>=5; extra == "dev"
+Requires-Dist: mypy>=1.11; extra == "dev"
+Requires-Dist: ruff>=0.5; extra == "dev"
+Requires-Dist: pre-commit>=3.7; extra == "dev"
+Dynamic: license-file
+
+# 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 a `Trainer` class for running epoch-based training. It supports training either from a `DataLoader` or directly from tensors, and records basic diagnostics such as training 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](https://github.com/jankazil/torch-tk/blob/main/notebooks/HowTo.ipynb) Jupyter notebook. 
+
+## Installation
+
+### pip
+```bash
+pip install torch-tk
+```
+
+### conda / mamba
+```bash
+mamba install -c jan.kazil -c conda-forge torch-tk
+```
+
+## 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
+
+- `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
+
+## Public API
+
+### 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 the `Trainer` class for epoch-based training and simple training diagnostics.
+
+- `Trainer(model, optimizer, loss_function, epoch=0)`: Initialize trainer state.
+- `Trainer.train_with_dataloader(data_loader, num_epochs, epoch_diag_step=1, verbose=True)`: Train from a `DataLoader`.
+- `Trainer.train_with_data(x_train, y_train, bs, num_epochs, epoch_diag_step=1, verbose=True, shuffle=False)`: Train from in-memory tensors.
+- `Trainer.plot_loss(...)`: Plot recorded diagnostic loss versus epoch.
+- `Trainer.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)`: 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` 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` container for sample-resolved loss diagnostics and analysis.
+
+- `Diagnostics.from_data_loader(...)`: Build diagnostics from a model evaluated on a `DataLoader`.
+- `Diagnostics.from_data(...)`: Build diagnostics from in-memory tensors.
+- `Diagnostics.from_netcdf(path)`: Restore diagnostics 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_diagnostics(diagnostics, plot_file=None, title=None, font_factor=1.5, figsize=(9, 6), xlim=None, ylim=None, loss_name='sqrt(loss)', pdf_bin_n=100, dpdlog10=False, show_plot=True, verbose=True)`: Plot kernel-density estimates of square-root per-sample loss distributions across one or more diagnostics objects and epochs.
+
+## 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 diagnostic plotting utility requires strictly positive, non-negative loss values because it plots the square root of loss on a logarithmic axis.
+- 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 - jan.kazil.dev@gmail.com - [jankazil.com](https://jankazil.com)
+
+## License
+
+BSD-3-Clause

torch_tk-1.0.8/README.md

@@ -0,0 +1,162 @@
+# 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 a `Trainer` class for running epoch-based training. It supports training either from a `DataLoader` or directly from tensors, and records basic diagnostics such as training 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](https://github.com/jankazil/torch-tk/blob/main/notebooks/HowTo.ipynb) Jupyter notebook. 
+
+## Installation
+
+### pip
+```bash
+pip install torch-tk
+```
+
+### conda / mamba
+```bash
+mamba install -c jan.kazil -c conda-forge torch-tk
+```
+
+## 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
+
+- `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
+
+## Public API
+
+### 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 the `Trainer` class for epoch-based training and simple training diagnostics.
+
+- `Trainer(model, optimizer, loss_function, epoch=0)`: Initialize trainer state.
+- `Trainer.train_with_dataloader(data_loader, num_epochs, epoch_diag_step=1, verbose=True)`: Train from a `DataLoader`.
+- `Trainer.train_with_data(x_train, y_train, bs, num_epochs, epoch_diag_step=1, verbose=True, shuffle=False)`: Train from in-memory tensors.
+- `Trainer.plot_loss(...)`: Plot recorded diagnostic loss versus epoch.
+- `Trainer.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)`: 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` 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` container for sample-resolved loss diagnostics and analysis.
+
+- `Diagnostics.from_data_loader(...)`: Build diagnostics from a model evaluated on a `DataLoader`.
+- `Diagnostics.from_data(...)`: Build diagnostics from in-memory tensors.
+- `Diagnostics.from_netcdf(path)`: Restore diagnostics 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_diagnostics(diagnostics, plot_file=None, title=None, font_factor=1.5, figsize=(9, 6), xlim=None, ylim=None, loss_name='sqrt(loss)', pdf_bin_n=100, dpdlog10=False, show_plot=True, verbose=True)`: Plot kernel-density estimates of square-root per-sample loss distributions across one or more diagnostics objects and epochs.
+
+## 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 diagnostic plotting utility requires strictly positive, non-negative loss values because it plots the square root of loss on a logarithmic axis.
+- 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 - jan.kazil.dev@gmail.com - [jankazil.com](https://jankazil.com)
+
+## License
+
+BSD-3-Clause

torch_tk-1.0.8/pyproject.toml

@@ -0,0 +1,98 @@
+[build-system]
+requires = ["setuptools>=77", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "torch-tk"
+version = "1.0.8"
+description = "Streamlines training, checkpoint management, and diagnostics of PyTorch models."
+readme = "README.md"
+requires-python = ">=3.12"
+license = "BSD-3-Clause"
+license-files = ["LICENSE"]
+authors = [
+    { name = "Jan Kazil", email = "jan.kazil.dev@gmail.com" }
+]
+dependencies = [
+  "torch",
+  "numpy",
+  "xarray",
+  "matplotlib",
+  "scipy",
+]
+classifiers = [
+    "Development Status :: 4 - Beta",
+    "Intended Audience :: Developers",
+    "Intended Audience :: Science/Research",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.12",
+    "Topic :: Scientific/Engineering",
+    "Topic :: Software Development :: Libraries :: Python Modules",
+]
+keywords = ["pytorch", "training", "checkpointing", "diagnostics", "deep-learning"]
+
+[project.urls]
+Homepage = "https://github.com/jankazil/torch-tk"
+Repository = "https://github.com/jankazil/torch-tk"
+Issues = "https://github.com/jankazil/torch-tk/issues"
+
+[project.optional-dependencies]
+dev = [
+  "build>=1.2",
+  "twine>=5.1",
+  "pytest>=8",
+  "pytest-cov>=5",
+  "mypy>=1.11",
+  "ruff>=0.5",
+  "pre-commit>=3.7",
+]
+
+[tool.setuptools]
+package-dir = {"" = "src"}
+
+[tool.setuptools.packages.find]
+where = ["src"]
+
+[tool.ruff]
+line-length = 128
+target-version = "py312"
+extend-exclude = [
+    "__init__.py",
+    "dist",
+    "build",
+    "data",
+    "demos",
+    "docs",
+    "experiments",
+    "notebooks",
+]
+
+[tool.ruff.lint]
+select = ["E", "F", "I", "UP", "B", "SIM"]
+ignore = ["E501", "SIM108"]
+
+[tool.ruff.format]
+quote-style = "preserve"
+
+#
+# MyPy is disabled because static analysis cannot account for all
+# dynamic runtime behaviors, mypy may report false positives which
+# do no reflect actual runtime issues.
+#
+#[tool.mypy]
+#python_version = "3.12"
+#warn_unused_configs = true
+#disallow_untyped_defs = true
+#disallow_incomplete_defs = true
+#no_implicit_optional = true
+#check_untyped_defs = true
+#strict_optional = true
+#pretty = true
+#namespace_packages = true
+#mypy_path = "src"
+#files = ["src/torch_tk", "scripts"]
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+python_files = ["test_*.py"]
+python_functions = ["test_*"]

torch_tk-1.0.8/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

torch_tk-1.0.8/src/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-1.0.8/src/torch_tk/checkpoints/__init__.py

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

torch_tk-1.0.8/src/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