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()})