Model Deployment: ONNX, TensorRT, Quantization, and Pruning

🎯 Learning Objectives

• Export PyTorch models to ONNX for cross-framework deployment
• Optimize models with TensorRT for GPU inference
• Apply quantization (INT8, FP16) to reduce model size and latency
• Prune models to remove redundant parameters
• Serialize models for production serving

📋 Prerequisites

• PyTorch basics (Week 1): Model definition, forward pass
• Profiling (Week 7): Understanding bottlenecks

1. 📖 Core Content

1.1 Intuition: From Research to Production

1. Faster inference: Reduce latency from 50ms to 5ms
2. Smaller size: Reduce from 500MB to 50MB for mobile
3. Cross-platform: Run on Python, C++, Java, JavaScript, mobile
4. Hardware optimization: Maximize GPU/CPU/TPU utilization The deployment pipeline: PyTorch → ONNX → TensorRT/OpenVINO → Production

1.2 ONNX Export

python
Run Inline
# runnable
import torch
import torch.nn as nn
class SimpleModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = nn.Conv2d(3, 16, 3)
        self.relu = nn.ReLU()
        self.pool = nn.AdaptiveAvgPool2d(1)
        self.fc = nn.Linear(16, 10)
    def forward(self, x):
        x = self.relu(self.conv(x))
        x = self.pool(x).squeeze(-1).squeeze(-1)
        x = self.fc(x)
        return x
model = SimpleModel()
model.eval()
# Export to ONNX
dummy_input = torch.randn(1, 3, 32, 32)
torch.onnx.export(
    model,
    dummy_input,
    "model.onnx",
    input_names=['input'],
    output_names=['output'],
    dynamic_axes={
        'input': {0: 'batch_size'},  # Variable batch size
        'output': {0: 'batch_size'}
    },
    opset_version=17
)
print("✅ Model exported to model.onnx")
# Verify ONNX
import onnx
onnx_model = onnx.load("model.onnx")
onnx.checker.check_model(onnx_model)
print(f"✅ ONNX model verified: {onnx_model.graph.node[0].op_type}")

1.3 TensorRT Optimization

python
Run Inline
# runnable
import torch
import tensorrt as trt
# Convert ONNX to TensorRT engine
TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
builder = trt.Builder(TRT_LOGGER)
network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
parser = trt.OnnxParser(network, TRT_LOGGER)
with open("model.onnx", "rb") as f:
    parser.parse(f.read())
config = builder.create_builder_config()
config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 20)  # 1MB
# FP16 mode (if supported)
if builder.platform_has_fast_fp16:
    config.set_flag(trt.BuilderFlag.FP16)
    print("Using FP16 mode")
# Build engine
serialized_engine = builder.build_serialized_network(network, config)
with open("model.engine", "wb") as f:
    f.write(serialized_engine)
print("✅ TensorRT engine built")

1.4 Quantization

python
Run Inline
# runnable
import torch
# Post-training quantization
model = SimpleModel().eval()
# Prepare for quantization
model.qconfig = torch.quantization.get_default_qconfig('fbgemm')  # For CPU
quantized_model = torch.quantization.quantize_dynamic(
    model,
    {nn.Linear, nn.Conv2d},  # Quantize these layers
    dtype=torch.qint8
)
# Measure size reduction
import sys
def get_model_size(model):
    torch.save(model.state_dict(), "temp.pth")
    size = sys.getsizeof(open("temp.pth", "rb").read())
    return size
original_size = get_model_size(model)
quantized_size = get_model_size(quantized_model)
print(f"Original: {original_size:,} bytes")
print(f"Quantized: {quantized_size:,} bytes")
print(f"Compression: {original_size/quantized_size:.1f}×")

1.5 Model Pruning

python
Run Inline
# runnable
import torch.nn.utils.prune as prune
model = SimpleModel()
# Apply L1 unstructured pruning (50% of conv1 weights)
prune.l1_unstructured(model.conv, name='weight', amount=0.5)
# Verify pruning
print(f"Remaining non-zero weights: {torch.count_nonzero(model.conv.weight).item()}")
print(f"Total weights: {model.conv.weight.numel()}")
print(f"Sparsity: {1 - torch.count_nonzero(model.conv.weight).item()/model.conv.weight.numel():.1%}")
# Make pruning permanent
prune.remove(model.conv, 'weight')

1.6 Why This Matters

2. 📐 Key Formulas / Concepts

3. ⚠️ Common Pitfalls

Pitfall 1: Dynamic Control Flow in ONNX

Pitfall 2: Quantization Without Calibration

4. 📝 Practice Questions

Q1: A model takes 150ms for FP32 inference on CPU and 45ms for INT8 inference. The INT8 model has 98% of FP32 accuracy. Should you deploy INT8?

Yes, if the latency improvement is needed. The 3.3× speedup and minimal accuracy loss (2%) make INT8 well worth it for most production applications. Check:

1. Accuracy requirement: If 98% is above the business threshold, deploy.
2. Latency budget: If 150ms is too slow (e.g., real-time requirement < 100ms), INT8 is the solution.
3. Target hardware: INT8 requires hardware support (VNNI on modern CPUs, Tensor Cores on GPUs). Q2: Your exported ONNX model runs slower than the original PyTorch model. Why?
4. PyTorch uses graph optimization automatically (JIT). ONNX doesn't apply these optimizations unless explicitly configured.
5. Op fusion: PyTorch fuses operations internally. ONNX keeps them separate unless the inference engine fuses them.
6. TensorRT optimization: Running ONNX without TensorRT optimization leaves performance on the table. Always optimize with TensorRT, OpenVINO, or ONNX Runtime.
7. Unsupported ops: Some PyTorch ops may fall back to a reference implementation in the inference engine.

5. 🔗 Cross-References

• Previous: Profiling & Optimization (Week 7)
• Next: Experiment Tracking (Week 9)
• Related: DL-CV Cross-course
• External: NVIDIA TensorRT Documentation Join Discord PreviousProfiling & OptimizationNextExperiment Tracking