Deploying Vision Models

Build the mental model

ONNX (Open Neural Network Exchange) is an intermediate format that captures a model's computation graph and weights independently of the framework that trained it. A model trained in PyTorch can then run on ONNX Runtime, TensorRT, mobile inference engines on Android or iOS, or even in a web browser via onnxruntime-web, all without needing a full PyTorch installation. torch.onnx.export traces — or symbolically captures — the model's forward pass given a representative dummy input, recording every operation as a node in the exported graph.

Connect it to a real scenario

The Tutorial Platform's upload pipeline runs a vision model that auto-generates captions for diagram images, but bundling the entire PyTorch runtime into the serverless function that processes uploads makes cold-start time unacceptably slow. The platform team exports that model to ONNX and swaps the deployment dependency for ONNX Runtime's much lighter package, cutting the serverless function's cold-start latency significantly.

ONNX is a framework-independent format that lets a trained model run on many inference runtimes, including edge devices.
torch.onnx.export uses a dummy input to capture the model's forward pass as a graph and save it as a .onnx file.
Quantization — converting weights and activations from float32 to lower precision like int8 — shrinks model size and speeds up inference on edge and CPU hardware, typically at a small accuracy cost.

Try the working example

python

import torch
import torch.nn as nn

torch.manual_seed(0)


class TinyClassifier(nn.Module):
    def __init__(self, num_classes=10):
        super().__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(3, 8, kernel_size=3, stride=2, padding=1),  # 32x32 -> 16x16
            nn.ReLU(),
            nn.AdaptiveAvgPool2d(1),
        )
        self.classifier = nn.Linear(8, num_classes)

    def forward(self, x):
        features = self.conv(x).flatten(1)
        return self.classifier(features)


model = TinyClassifier()
model.eval()  # disable dropout/batchnorm training behavior before export

dummy_input = torch.randn(1, 3, 32, 32)  # one fake RGB image, batch size 1

onnx_path = "tiny_classifier.onnx"
# Requires the onnx and onnxscript packages in addition to torch/torchvision:
#   pip install onnx onnxscript
torch.onnx.export(
    model,
    dummy_input,
    onnx_path,
    input_names=["image"],
    output_names=["class_scores"],
    dynamic_axes={"image": {0: "batch_size"}, "class_scores": {0: "batch_size"}},
    opset_version=18,
    verbose=False,
)

print(f"Exported model to {onnx_path}")
print("This ONNX file can now run on ONNX Runtime, mobile, or edge inference engines")
print("Input shape:", dummy_input.shape)
print("Output shape:", model(dummy_input).shape)

You should see

The export creates a tiny_classifier.onnx file in the working directory (the exporter may also print some internal diagnostic logging to the console during the export itself). The explicit prints then show Exported model to tiny_classifier.onnx, a line describing the file's intended use, Input shape: torch.Size([1, 3, 32, 32]), and Output shape: torch.Size([1, 10]) — a score vector over 10 classes.

5-minute try-it

Modify the code to change dummy_input from a batch of 1 to torch.randn(4, 3, 32, 32); since dynamic_axes is already configured, the export should still succeed. Print how the exported model's output shape changes for a batch of 4, and think through what would go wrong if dynamic_axes had been omitted entirely.

One important caution

Exporting with a fixed batch size of 1 and no dynamic_axes means that when deployment code later sends a batch of 8 images, the ONNX runtime raises a shape mismatch error.

Forgetting to call model.eval() before export leaves dropout and batchnorm layers in training mode, baking randomness or batch-dependent statistics into the traced graph — so the exported model's outputs for single-image inference can disagree with the original trained model.

PyTorch Docs — torch.onnx — Computer Vision