Complete PyTorch Series — Module 12 of 12

Building and training state-of-the-art models is only half the battle. In production, models must serve predictions under stringent Service Level Agreements (SLAs)—often requiring response times under 20 milliseconds at thousands of requests per second. In this final capstone module, you will master model serialization via TorchScript, framework-agnostic export with ONNX, and deploy an asynchronous REST API using FastAPI.

Exporting Models to TorchScript (C++ Independent Runtime)

Python has the Global Interpreter Lock (GIL) and significant runtime overhead. TorchScript compiles your PyTorch model into an intermediate representation (IR) graph that can run directly inside standalone C++ production servers without Python installed:

Method 1: Tracing with torch.jit.trace

Passes an example input tensor through the model and records the exact operations executed. Ideal for vision models and architectures with fixed control flow:

import torch
import torch.nn as nn

class Classifier(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(nn.Linear(10, 32), nn.ReLU(), nn.Linear(32, 2))
    def forward(self, x):
        return self.net(x)

model = Classifier().eval()
dummy_input = torch.randn(1, 10)

# Trace the model
traced_model = torch.jit.trace(model, dummy_input)

# Save standalone binary
traced_model.save("model_traced.pt")
print("TorchScript model saved as model_traced.pt")

# Load and execute in pure TorchScript (even in C++!)
loaded_jit = torch.jit.load("model_traced.pt")
output = loaded_jit(dummy_input)
print("TorchScript output:", output)

Exporting to ONNX (Open Neural Network Exchange)

ONNX is an open format built to represent machine learning models across platforms. Once exported to ONNX, your model can run on ultra-optimized inference engines like ONNX Runtime, NVIDIA TensorRT, or mobile edges (iOS CoreML, Android NNAPI):

import torch

# Export to ONNX with dynamic batch dimension
torch.onnx.export(
    model,                          # model being run
    dummy_input,                    # model input (or a tuple for multiple inputs)
    "classifier.onnx",              # where to save the model
    export_params=True,             # store the trained parameter weights inside the model file
    opset_version=17,               # the ONNX version to export the model to
    do_constant_folding=True,       # optimize constant foldings
    input_names=['input_features'], # the model's input names
    output_names=['class_logits'],  # the model's output names
    dynamic_axes={
        'input_features': {0: 'batch_size'}, # variable length axes
        'class_logits': {0: 'batch_size'}
    }
)
print("Model successfully exported to classifier.onnx with dynamic batch size!")

Running Inference with ONNX Runtime

import onnxruntime as ort
import numpy as np

# Initialize ONNX Runtime session
session = ort.InferenceSession("classifier.onnx", providers=['CPUExecutionProvider'])

# Prepare input as raw NumPy array
input_data = np.random.randn(4, 10).astype(np.float32) # Batch of 4

# Run inference
ort_inputs = {session.get_inputs()[0].name: input_data}
ort_outputs = session.run(None, ort_inputs)

print("ONNX Runtime predictions shape:", ort_outputs[0].shape) # (4, 2)

Building a High-Performance Serving API with FastAPI

FastAPI provides an asynchronous, ASGI-based web server with automatic Swagger documentation, Pydantic data validation, and ultra-low latency.

Create serve.py:

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from typing import List
import torch
import torch.nn as nn
from contextlib import asynccontextmanager

# Define Pydantic request/response schemas
class PredictRequest(BaseModel):
    features: List[List[float]] # Nested list: batch of samples

class PredictResponse(BaseModel):
    predicted_classes: List[int]
    probabilities: List[List[float]]

# Global model container
ml_models = {}

# Lifespan context manager: loads model into memory once on startup
@asynccontextmanager
async def lifespan(app: FastAPI):
    # Load model weights on startup
    print("Loading PyTorch model into memory...")
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    
    # Simple model architecture
    model = nn.Sequential(nn.Linear(10, 32), nn.ReLU(), nn.Linear(32, 2)).to(device)
    model.eval()
    
    ml_models["model"] = model
    ml_models["device"] = device
    print("Model loaded and ready for inference.")
    yield
    # Cleanup on shutdown
    ml_models.clear()

app = FastAPI(title="PyTorch Production API", lifespan=lifespan)

@app.get("/health")
def health_check():
    return {"status": "healthy", "device": str(ml_models["device"])}

@app.post("/predict", response_model=PredictResponse)
def predict(request: PredictRequest):
    try:
        device = ml_models["device"]
        model = ml_models["model"]

        # Convert input payload to PyTorch tensor
        input_tensor = torch.tensor(request.features, dtype=torch.float32, device=device)

        if input_tensor.ndim != 2 or input_tensor.shape[1] != 10:
            raise ValueError("Input features must have shape (Batch_Size, 10)")

        # Run inference in eval mode without Autograd overhead
        with torch.no_grad():
            logits = model(input_tensor)
            probs = torch.softmax(logits, dim=-1)
            predicted_classes = torch.argmax(probs, dim=-1).cpu().tolist()
            prob_list = probs.cpu().tolist()

        return PredictResponse(
            predicted_classes=predicted_classes,
            probabilities=prob_list
        )
    except Exception as e:
        raise HTTPException(status_code=400, detail=str(e))

Production Deployment Checklist

• Disable Gradient Calculation: Ensure every prediction endpoint runs inside with torch.no_grad(): or uses torch.inference_mode().
• Lock Evaluation State: Always call model.eval() immediately after loading weights.
• Tune CPU Threads: If serving on CPU nodes, configure torch.set_num_threads(4) to prevent thread contention.
• Containerize: Package your FastAPI application using a minimal Docker container based on python:3.11-slim or NVIDIA CUDA runtime images.
• Load Balancer & Workers: Run your server using uvicorn serve:app --host 0.0.0.0 --port 8000 --workers 4 behind Nginx or an AWS/GCP Application Load Balancer.

Module 12 Recap

• TorchScript serializes models into an independent representation executable in C++ without Python.
• ONNX creates a universal open format runnable on ONNX Runtime and TensorRT for maximal throughput.
• FastAPI provides asynchronous, typed endpoints for high-throughput model serving.
• Load models into memory once at application startup using lifespan events.