torch-tk

torch-tk provides a set of utilities to simplify the PyTorch workflow, focusing on making models and optimizers self-describing for easier serialization and reconstruction. Key capabilities include: - **Model & Optimizer Management**: Offers a `Model` base class and optimizer wrappers (e.g., `SGD`, `Adam`) that store constructor arguments and state, allowing them to be saved to and restored from files automatically. - **Checkpointing**: A `CheckPointManager` that handles the simultaneous saving and loading of both the model and optimizer states. - **Training Utilities**: Includes a standard `Trainer` for typical datasets and a `MaskedTrainer` specifically designed for targets containing missing values (NaNs). Both support training via `DataLoader` or direct tensors. - **Diagnostics**: A `Diagnostics` class for computing and storing sample-resolved loss information. It supports saving diagnostics to netCDF files and provides plotting utilities for loss probability distributions (KDE and histograms). The toolkit is designed to preserve the standard PyTorch interface while automating the boilerplate associated with training loops and state persistence.