Principle:LaurentMazare Tch rs Trainable JIT Loading
| Knowledge Sources | |
|---|---|
| Domains | Model_Training, Interoperability |
| Last Updated | 2026-02-08 14:00 GMT |
Overview
Pattern for loading a TorchScript model in Rust with its parameters registered in a VarStore, enabling fine-tuning and continued training.
Description
Trainable JIT loading extends standard TorchScript loading by extracting all named parameters from the loaded module and registering them in a tch-rs VarStore. This enables the parameters to be tracked by optimizers for gradient-based training. It bridges the gap between Python model definition and Rust training, allowing a model exported from Python to be fine-tuned entirely in Rust.
Usage
Use when you need to train or fine-tune a TorchScript model in Rust rather than just running inference. The model's parameters become trainable through the VarStore optimizer system.
Theoretical Basis
TrainableCModule loading:
1. Load CModule from .pt file onto VarStore's device
2. Extract named_parameters() from the C++ module
3. For each (name, tensor):
- Register in VarStore with requires_grad preserved
4. Now optimizers can train these parameters
Inference-only vs Trainable:
CModule::load → Parameters in C++ land only
TrainableCModule::load → Parameters also in VarStore → Optimizer can update them
Related Pages
Implemented By
Page Connections
Double-click a node to navigate. Hold to expand connections.
Principle
Implementation
Heuristic
Environment