Workflow:LaurentMazare Tch rs JIT Model Inference
| Knowledge Sources | |
|---|---|
| Domains | Deep_Learning, Model_Deployment, TorchScript, Rust_ML |
| Last Updated | 2026-02-08 13:00 GMT |
Overview
End-to-end process for exporting a trained PyTorch model from Python as a TorchScript module and loading it for inference in Rust using the tch-rs CModule API.
Description
This workflow bridges the Python PyTorch training ecosystem with Rust deployment. A model trained in Python is exported via TorchScript (tracing or scripting), producing a serialized .pt file. The tch-rs CModule API loads this file in Rust and executes the forward pass without any Python dependency at runtime. This enables high-performance, safe deployment of PyTorch models in Rust applications, including support for training JIT models in Rust via TrainableCModule.
Usage
Execute this workflow when you have a model trained in Python and want to deploy it in a Rust application for inference (or continued training). This is the recommended approach for production deployment where Python overhead is undesirable, or when integrating PyTorch models into existing Rust systems.
Execution Steps
Step 1: Export the model from Python
In Python, use torch.jit.trace or torch.jit.script to convert a trained PyTorch model into a TorchScript module, then save it to a .pt file. Tracing records operations by running a sample input through the model; scripting compiles the Python code directly.
Key considerations:
- torch.jit.trace requires a representative sample input tensor with correct shape and dtype
- Traced models may not capture data-dependent control flow (use torch.jit.script for that)
- The exported .pt file is self-contained with both architecture and weights
Step 2: Load the TorchScript module in Rust
Use CModule::load to deserialize the .pt file into a callable module. For inference-only use, CModule provides forward_ts (tensor list input) and forward_is (IValue input). For training, use TrainableCModule::load which registers parameters with a VarStore.
Key considerations:
- CModule::load takes a file path and returns a module ready for inference
- CModule::load_on_device allows specifying the target device (CPU/CUDA) at load time
- TrainableCModule::load takes a VarStore root path to enable gradient tracking
Step 3: Preprocess the input
Prepare the input tensor in the same format the model expects. For vision models, this typically means loading an image, resizing to the expected dimensions, normalizing, and adding a batch dimension. The preprocessing must match exactly what was used during Python training.
Key considerations:
- Input tensor shape, dtype, and normalization must match the Python training pipeline
- The imagenet module provides helpers for standard vision preprocessing
- For non-vision models, construct the input tensor from raw data as appropriate
Step 4: Run the forward pass
Apply the loaded module to the input tensor using the Module trait's forward method or the CModule-specific forward_ts/forward_is methods. The output is a Tensor (or IValue for complex outputs) containing the model's predictions.
What happens:
- The apply method calls forward_ts internally with a single-tensor list
- For models with multiple inputs, use forward_ts with a Vec of tensors
- For models returning complex types (tuples, dicts), use forward_is with IValue
Step 5: Post-process results
Convert the output tensor to human-readable results. For classification, apply softmax and extract top-k class indices. For other tasks, decode the output according to the model's output specification.
Key considerations:
- Softmax converts logits to probabilities
- imagenet::top provides top-k extraction with class name lookup for ImageNet models
- For custom models, implement task-specific output decoding