Python / Python Deep Learning and Neural Networks Interview Questions
How do you export a PyTorch model for production deployment using TorchScript or ONNX?
Research-time PyTorch models depend on Python's interpreter and PyTorch's eager execution mode — both are too slow and have too many dependencies for production deployment. Two standard serialisation formats allow deploying PyTorch models without Python: TorchScript (PyTorch-native, supports dynamic shapes better) and ONNX (framework-agnostic, runs on TensorRT, OpenVINO, CoreML, ONNX Runtime across many hardware targets).
TorchScript compiles a PyTorch model into an intermediate representation (IR) that can run in C++ via LibTorch, without any Python dependency. It is created either via torch.jit.trace (records operations from a concrete example — doesn't handle data-dependent control flow) or torch.jit.script (analyzes Python source — handles control flow but requires type annotations and a subset of Python). ONNX export traces the model similarly and serialises it to the ONNX protobuf format, which can then be run on any ONNX-compatible runtime.
import torch import torch.nn as nn model = nn.Sequential( nn.Linear(10, 64), nn.ReLU(), nn.Linear(64, 5) ).eval() example_input = torch.randn(1, 10) # âââ TorchScript: tracing âââââââââââââââââââââââââââââââââââââââââââ traced = torch.jit.trace(model, example_input) traced.save('model_traced.pt') # Load and run without original Python class: loaded = torch.jit.load('model_traced.pt') output = loaded(torch.randn(4, 10)) # âââ TorchScript: scripting (handles if/for in forward) ââââââââââââ @torch.jit.script def activation(x: torch.Tensor) -> torch.Tensor: if x.sum() > 0: return torch.relu(x) return torch.tanh(x) # âââ ONNX export ââââââââââââââââââââââââââââââââââââââââââââââââââââ torch.onnx.export( model, example_input, 'model.onnx', input_names=['features'], output_names=['logits'], dynamic_axes={'features': {0: 'batch_size'}, # variable batch 'logits': {0: 'batch_size'}}, opset_version=17, ) # Validate exported ONNX model import onnx, onnxruntime as ort onnx.checker.check_model('model.onnx') sess = ort.InferenceSession('model.onnx') result = sess.run(None, {'features': example_input.numpy()})
More Related questions...