Deploying PyTorch Models: Export, ONNX, Quantization & Serving

Q: Should I use TorchScript, torch.export, or ONNX?

For new projects, use torch.export when staying in the PyTorch ecosystem (AOTInductor, ExecuTorch), and ONNX when targeting ONNX Runtime, TensorRT, or other runtimes and languages. TorchScript still works, but it's in maintenance mode. If you're serving from Python anyway, eager mode with torch.compile may be all you need.

Deploying PyTorch Models

Training produces a model; deployment makes it useful. A deployed PyTorch model must load reliably, run fast enough and cheaply enough on target hardware, handle concurrent requests, and be versioned, monitored, and updated like any other service. The path from notebook to production involves choices at several layers: format (eager Python, torch.export, ONNX), optimization (compilation, quantization), serving (a web app, Triton, vLLM, or a managed endpoint), and hardware (GPU, CPU, mobile, or edge).

The right choice depends on the model and workload. A small classifier behind a FastAPI endpoint on CPUs, a vision model on Triton with dynamic batching, and a large language model on vLLM have very different needs.

TL;DR

Save state_dicts and load with weights_only=True (or safetensors). Never unpickle untrusted model files.
For portable, Python-free graphs, use torch.export (the modern path), or export to ONNX for ONNX Runtime and TensorRT.
Speed up inference with torch.compile, half precision, quantization (int8, int4), and batching.
Serve with a simple FastAPI app for light workloads, Triton or TorchServe for high-throughput multi-model serving, and vLLM/SGLang for LLMs.
Use dynamic batching on GPUs to raise throughput, and always put model.eval() plus inference_mode() in the inference path.
Treat models as versioned artifacts: registry, reproducible preprocessing, monitoring, and rollback.

Quick Example

A minimal but production-shaped inference service with FastAPI:

Exporting the same model for runtimes outside Python:

Core Concepts

Saving and Loading Safely

Save model.state_dict(), plus the config needed to rebuild the architecture.
Load with torch.load(..., weights_only=True). Full unpickling can execute arbitrary code (see insecure deserialization), and recent PyTorch versions default to weights-only loading.
safetensors is a pickle-free format for weights, widely used on Hugging Face.
Version models with their training data version, code commit, metrics, and preprocessing in a model registry (MLflow, W&B, SageMaker, Vertex). See MLOps.

Export Formats

Export requires models whose control flow and shapes can be captured. Mark dynamic dimensions (batch size, sequence length) explicitly.

Inference Optimization

model.eval() + torch.inference_mode(): always.
torch.compile or AOTInductor: fused kernels and lower overhead.
Reduced precision: fp16 or bf16 on GPUs, often with negligible accuracy loss.
Quantization: int8 (dynamic or static) for CPUs and GPUs; 4-bit weight quantization (GPTQ, AWQ) for LLMs; quantization-aware training when post-training quantization hurts accuracy too much.
TensorRT (via Torch-TensorRT or ONNX) for maximum NVIDIA GPU performance.
Distillation and pruning: smaller models for tight latency or cost budgets.

Always validate accuracy after optimization against a held-out set.

Serving Options

Batching and Hardware

GPUs are most efficient on batches. Dynamic batching groups concurrent requests within a few milliseconds into one forward pass, raising throughput dramatically at small latency cost. For small models with low traffic, CPU inference (with quantization and ONNX Runtime) is often cheaper and simpler than GPUs. Right-size the hardware from load tests measuring p50, p95, and p99 latency against throughput.

Edge and Mobile

ExecuTorch runs exported PyTorch models on mobile and embedded devices with delegates for NPUs and GPUs. ONNX Runtime Mobile and Core ML conversion are alternatives. Quantization and small architectures are essential there. See embedded IoT.

Best Practices

Ship Preprocessing With the Model

Training-serving skew, where inference uses different tokenization, normalization, or feature logic, is a leading cause of silent accuracy loss. Package preprocessing code and parameters with the model version, and test them together.

Warm Up and Load Once

Load models at startup (lifespan hooks), move them to the target device, and run warm-up inference to trigger compilation and allocate memory, before marking the service ready.

Monitor Model Behavior, Not Just Uptime

Track latency, throughput, errors, GPU utilization, and model-level signals: prediction distributions, confidence, input drift, and outcome metrics where labels arrive later. Roll out new versions with canaries or shadow traffic. See MLOps.

Isolate Heavy Inference From Web Workers

CPU-bound inference inside an async web server blocks the event loop. Run inference in dedicated worker processes or model servers, and call them from lightweight API services. See FastAPI deployment.

Common Mistakes

Serving in Training Mode

Forgetting model.eval() leaves dropout active and batch norm using batch statistics, producing noisy, inconsistent predictions.

Loading Untrusted Checkpoints With Pickle

Per-Request Model Loading

Loading weights inside the request handler adds seconds per request and exhausts memory. Load once and reuse.

FAQ

How do I deploy a PyTorch model to production?

Save the trained weights and preprocessing, choose a serving approach (a Python web service for simple cases, a model server like Triton for scale, or an LLM server like vLLM for language models), optimize inference (eval mode, reduced precision, compilation or quantization), package it in a container, and deploy with health checks, autoscaling, monitoring, and versioned rollouts.

Should I use TorchScript, torch.export, or ONNX?

For new projects, use torch.export when staying in the PyTorch ecosystem (AOTInductor, ExecuTorch), and ONNX when targeting ONNX Runtime, TensorRT, or other runtimes and languages. TorchScript still works, but it's in maintenance mode. If you're serving from Python anyway, eager mode with torch.compile may be all you need.

How much does quantization reduce accuracy?

Often very little: int8 post-training quantization typically costs under 1% on many vision and NLP tasks, and 4-bit weight quantization of LLMs often retains most quality. Results vary by model and task, so always evaluate. Quantization-aware training can recover losses.

Do I need a GPU for inference?

Not always. Small and medium models, especially quantized and exported to ONNX Runtime, can serve efficiently on CPUs with lower cost and simpler operations. Large models, high throughput, or strict latency requirements usually justify GPUs, ideally with batching to keep them busy.