PyTorch Model Deployment
Deployment takes a trained PyTorch model and makes it available for real-world predictions — as a web API, a mobile app, an embedded device, or a cloud service. This final topic covers the main deployment paths, each suited to a different production environment.
The Deployment Pipeline
Research / Training
│
▼ torch.save() / model.eval()
Trained Model (.pth or .pt)
│
├──▶ Flask / FastAPI REST API ← web server, any client
├──▶ TorchServe ← PyTorch's production model server
├──▶ TorchScript + C++ ← high-performance backend
├──▶ ONNX Runtime ← framework-agnostic inference
└──▶ PyTorch Mobile ← iOS / Android on-device
Step 1: Prepare the Model for Inference
Before deploying, always set the model to evaluation mode and wrap inference in torch.no_grad(). This disables dropout and batch norm's training behavior, reduces memory use, and speeds up inference.
import torch
import torch.nn as nn
# Load trained weights
model = MyModel()
model.load_state_dict(torch.load('best_model.pth', map_location='cpu'))
model.eval()
def predict(input_tensor):
with torch.no_grad():
output = model(input_tensor)
return outputDeployment Option 1: Flask REST API
Wrap the model in a Flask web server. Clients send POST requests with input data and receive predictions as JSON responses. This is the quickest path to serving a model over HTTP.
from flask import Flask, request, jsonify
import torch
app = Flask(__name__)
model = MyModel()
model.load_state_dict(torch.load('best_model.pth', map_location='cpu'))
model.eval()
@app.route('/predict', methods=['POST'])
def predict():
data = request.json['input'] # list of floats from client
x = torch.tensor(data, dtype=torch.float32).unsqueeze(0)
with torch.no_grad():
output = model(x)
return jsonify({'prediction': output.tolist()})
if __name__ == '__main__':
app.run(host='0.0.0.0', port=5000)
Client sends:
POST http://localhost:5000/predict
{"input": [1.2, 0.5, -0.3, 0.8]}
Server returns:
{"prediction": [[0.72, 0.19, 0.09]]}
Deployment Option 2: FastAPI (Modern REST API)
FastAPI is faster than Flask, has automatic input validation, and generates interactive API documentation automatically. It is the preferred framework for new ML API projects.
from fastapi import FastAPI
from pydantic import BaseModel
import torch
app = FastAPI()
model = MyModel()
model.load_state_dict(torch.load('best_model.pth', map_location='cpu'))
model.eval()
class PredictRequest(BaseModel):
input: list[float]
@app.post('/predict')
def predict(req: PredictRequest):
x = torch.tensor(req.input, dtype=torch.float32).unsqueeze(0)
with torch.no_grad():
out = model(x)
return {'prediction': out.tolist()}
# Run with: uvicorn app:app --host 0.0.0.0 --port 8000Deployment Option 3: TorchServe
TorchServe is PyTorch's official production model server. It handles batching, logging, multi-model serving, A/B testing, and REST endpoints automatically. It is the recommended path for production-grade serving at scale.
Workflow: 1. Archive model: torch-model-archiver --model-name mymodel ... 2. Start server: torchserve --start --model-store=model_store 3. Register: curl -X POST "http://localhost:8081/models?model_name=mymodel&url=mymodel.mar" 4. Predict: curl -X POST http://localhost:8080/predictions/mymodel -T input.json TorchServe handles: ✓ Multiple concurrent requests ✓ Request batching for GPU efficiency ✓ Health checks and metrics ✓ Model versioning and hot-swapping
Deployment Option 4: Docker Container
Packaging your model server in a Docker container makes deployment reproducible across any cloud provider or on-premise server.
# Dockerfile example FROM python:3.10-slim WORKDIR /app COPY requirements.txt . RUN pip install -r requirements.txt COPY best_model.pth . COPY app.py . EXPOSE 8000 CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "8000"]
Build and run: docker build -t my-model-api . docker run -p 8000:8000 my-model-api Deploy to cloud: AWS: push to ECR → deploy on ECS or EKS GCP: push to Artifact Registry → deploy on Cloud Run Azure: push to ACR → deploy on ACI or AKS
Deployment Option 5: PyTorch Mobile
For on-device inference without a network call — essential for privacy, offline use, or latency-sensitive apps — deploy directly to iOS or Android.
from torch.utils.mobile_optimizer import optimize_for_mobile
model.eval()
scripted = torch.jit.script(model)
optimized = optimize_for_mobile(scripted)
optimized._save_for_lite_interpreter('model_mobile.ptl')
# Bundle model_mobile.ptl with your iOS / Android appPreprocessing and Postprocessing in Deployment
A deployed model needs the same preprocessing applied during training — normalization, resizing, tokenization. Always package these steps with the model so the API doesn't receive raw inconsistent input.
from torchvision import transforms
preprocess = transforms.Compose([
transforms.Resize((224, 224)),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406],
[0.229, 0.224, 0.225])
])
def predict_image(pil_image):
x = preprocess(pil_image).unsqueeze(0)
with torch.no_grad():
logits = model(x)
probs = torch.softmax(logits, dim=1)
top1 = probs.argmax(dim=1).item()
return {'class_index': top1, 'confidence': probs[0, top1].item()}Deployment Checklist
Before deploying: ✓ model.eval() called ✓ torch.no_grad() wrapping inference ✓ Same preprocessing as training ✓ Input validation (type, shape, range) ✓ Error handling for bad inputs ✓ Tested on edge cases ✓ Logging and monitoring set up ✓ Version pinned (PyTorch, torchvision, etc.) ✓ GPU/CPU device consistent between training and serving
Choosing a Deployment Path
Use case Recommended approach ────────────────────────────────────────────────────────── Quick prototype API Flask or FastAPI Production API FastAPI + Docker or TorchServe High-throughput GPU serving TorchServe + batching C++ backend integration TorchScript → LibTorch Cross-framework compatibility ONNX + ONNX Runtime Mobile (iOS/Android) PyTorch Mobile (.ptl) Edge / microcontroller Quantize + TorchScript or ONNX
Summary
Model deployment takes a trained model into a production environment where it serves real predictions. Always call model.eval() and use torch.no_grad() at inference time. Flask and FastAPI serve models as REST APIs — FastAPI is the modern choice. TorchServe provides full production features including batching, monitoring, and versioning. Docker makes deployments reproducible across environments. PyTorch Mobile deploys models on-device to iOS and Android. ONNX export enables inference on hardware that has no PyTorch support. Always include the same preprocessing pipeline in your serving code that was used during training. Deployment is not an afterthought — a well-deployed model is what turns research into real-world impact.