Pytorch checkpoint serialization

pytorch-checkpoint-serialization

License: restricted — no clear open-source license detected for the underlying tool; verify licensing before commercial use or redistribution. <!-- asb-license-banner -->

Summary

Save and persist trained PyTorch graph neural network model states to disk after convergence or completion of training epochs. This skill ensures reproducibility and enables downstream transfer learning by capturing the learned weights, architecture parameters, and optimizer state.

When to use

Apply this skill when a PyTorch-based model (such as a graph neural network trained on molecular retention-time data) has completed training or reached convergence, and you need to preserve the model state for later inference, evaluation, or transfer learning on new datasets.

When NOT to use

• Model training has not converged or loss is still decreasing significantly
• Checkpoint storage space is unavailable or training is exploratory/temporary
• Model will not be reused across different datasets or sessions

Inputs

• Trained PyTorch model object (neural network instance)
• Optimizer state
• Training metadata (epoch count, loss history, validation metrics)

Outputs

• Model checkpoint file (.pt or .pth format) on disk
• Serialized state_dict containing learned weights and biases
• Optional metadata file recording training convergence metrics

How to apply

At the end of the training loop in the entry-point script (e.g., train_SMRT.py), call PyTorch's checkpoint save mechanism to write the model's state_dict and training metadata to disk. Typical implementations capture the model weights after each epoch or upon early stopping criteria, allowing monitoring via TorchMetrics loss and validation metrics to determine the optimal point for checkpoint serialization. The checkpoint should include the trained model architecture instantiated via PyTorch and PyG, enabling subsequent loading and fine-tuning during transfer learning workflows (e.g., train_transfer_FE.py).

Related tools

• PyTorch (Core framework providing model state serialization and checkpoint I/O)
• PyG (PyTorch Geometric) (Graph neural network layer definitions whose weights are captured in checkpoint)
• TorchMetrics (Monitors loss and validation metrics to determine optimal checkpoint timing)

Examples

python train_SMRT.py

Evaluation signals

• Checkpoint file exists on disk and is readable by torch.load()
• Loaded model state_dict has same keys and tensor shapes as trained model
• Transfer learning script (train_transfer_FE.py) successfully loads checkpoint without dimension mismatch errors
• Validation loss metrics recorded at checkpoint time show convergence (plateau or early stopping threshold met)
• Inference on held-out validation set using loaded checkpoint matches epoch metrics at save time

Limitations

• Checkpoint does not capture random seed state; reproducibility requires explicit seed logging
• Storage footprint scales with model size; large graph neural networks can produce multi-gigabyte checkpoints
• Cross-version PyTorch compatibility is not guaranteed; checkpoints saved in one PyTorch version may fail to load in another
• No automatic version control; multiple checkpoints require manual naming convention to track experiments

Evidence

• [intro] Save the trained model checkpoint to disk upon convergence or completion.: "Save the trained model checkpoint to disk upon convergence or completion."
• [intro] Monitor training via TorchMetrics for loss and validation metrics over epochs.: "Monitor training via TorchMetrics for loss and validation metrics over epochs."
• [intro] instantiate the graph neural network architecture (PyTorch + PyG) and train on the SMRT data.: "instantiate the graph neural network architecture (PyTorch + PyG) and train on the SMRT data."