pytorch-deployment

name	PyTorch Deployment
description	Comprehensive guide for deploying PyTorch models to production, covering export formats, optimization techniques, and deployment patterns.

PyTorch Deployment

Overview

PyTorch deployment involves exporting models to production-ready formats, optimizing for inference performance, and serving models through various deployment patterns. This skill covers TorchScript, ONNX export, TorchServe, model optimization techniques, inference optimization, FastAPI deployment, model versioning, A/B testing, monitoring, error handling, and performance benchmarking.

Prerequisites

• Understanding of PyTorch and deep learning models
• Knowledge of model training and evaluation
• Familiarity with web frameworks (FastAPI, Flask)
• Understanding of Docker and containerization
• Basic knowledge of cloud deployment concepts

Key Concepts

Model Export Formats

• TorchScript: PyTorch's intermediate representation for production deployment
• Tracing: Captures computation path from example inputs
• Scripting: Captures entire Python code including control flow
• ONNX: Open Neural Network Exchange for cross-framework compatibility
• TorchServe: PyTorch's model serving framework

Model Optimization

• Quantization: Reducing precision (FP32 → FP16/INT8) for efficiency
• Pruning: Removing less important weights from models
• Knowledge Distillation: Training smaller models from larger teacher models
• Model Compression: Techniques to reduce model size

Inference Optimization

• Batching: Processing multiple inputs together for efficiency
• GPU Utilization: Multi-GPU inference for throughput
• Mixed Precision: Using FP16 for faster computation
• Caching: Repeated computation results

Deployment Patterns

• FastAPI Server: REST API for model serving
• TorchServe: Production-ready model serving framework
• ONNX Runtime: High-performance inference engine
• Docker Deployment: Containerized model deployment

Implementation Guide

Model Export Formats

TorchScript

TorchScript is an intermediate representation of a PyTorch model that can be run in a high-performance environment such as C++.

Tracing vs Scripting:

import torch
import torch.nn as nn

class MyModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = nn.Conv2d(3, 64, 3)
        self.fc = nn.Linear(64 * 26 * 26, 10)

    def forward(self, x):
        x = self.conv(x)
        x = x.view(x.size(0), -1)
        return self.fc(x)

model = MyModel()
model.eval()

# Method 1: Tracing (captures actual computation path)
example_input = torch.randn(1, 3, 28, 28)
traced_model = torch.jit.trace(model, example_input)
traced_model.save("model_traced.pt")

# Method 2: Scripting (captures entire Python code)
scripted_model = torch.jit.script(model)
scripted_model.save("model_scripted.pt")

# Loading and inference
loaded_model = torch.jit.load("model_traced.pt")
output = loaded_model(example_input)

Handling Control Flow with Scripting:

class ConditionalModel(nn.Module):
    def forward(self, x):
        if x.sum() > 0:
            return x * 2
        else:
            return x / 2

# Use scripting for models with control flow
model = ConditionalModel()
scripted_model = torch.jit.script(model)

ONNX Export

Open Neural Network Exchange (ONNX) enables interoperability between different frameworks.

import torch
import torch.onnx

# Export to ONNX
model = MyModel()
model.eval()
dummy_input = torch.randn(1, 3, 28, 28)

torch.onnx.export(
    model,                      # Model to export
    dummy_input,                # Example input
    "model.onnx",               # Output file
    export_params=True,         # Store trained parameters
    opset_version=17,           # ONNX opset version
    do_constant_folding=True,   # Optimize constants
    input_names=['input'],      # Input names
    output_names=['output'],    # Output names
    dynamic_axes={              # Dynamic axes for variable batch size
        'input': {0: 'batch_size'},
        'output': {0: 'batch_size'}
    }
)

# Verify ONNX model
import onnx
onnx_model = onnx.load("model.onnx")
onnx.checker.check_model(onnx_model)

# Run inference with ONNX Runtime
import onnxruntime as ort

session = ort.InferenceSession("model.onnx")
input_name = session.get_inputs()[0].name
output_name = session.get_outputs()[0].name

outputs = session.run([output_name], {input_name: dummy_input.numpy()})

Custom ONNX Operators:

from torch.onnx import register_custom_op_symbolic

def custom_gsymbolic(g, input, alpha):
    return g.op("CustomOp", input, alpha_f=alpha)

register_custom_op_symbolic("aten::gelu", custom_gsymbolic, 17)

TorchServe

TorchServe is a flexible, easy-to-use tool for serving PyTorch models.

Installation:

pip install torchserve torch-model-archiver torch-workflow-archiver

Model Archiving:

# Create model handler (handler.py)
class ModelHandler:
    def __init__(self):
        self.model = None
        self.mapping = None
        self.device = None
        self.initialized = False

    def initialize(self, context):
        """Initialize model and load weights."""
        properties = context.system_properties
        self.device = torch.device("cuda:" + str(properties.get("gpu_id")) if torch.cuda.is_available() else "cpu")

        model_dir = properties.get("model_dir")
        model_pt_path = os.path.join(model_dir, "model.pth")

        self.model = torch.load(model_pt_path, map_location=self.device)
        self.model.eval()
        self.initialized = True

    def preprocess(self, requests):
        """Preprocess input data."""
        inputs = []
        for req in requests:
            data = req.get("data") or req.get("body")
            inputs.append(torch.tensor(data))
        return torch.stack(inputs)

    def ():
        
         torch.no_grad():
            output = .model(input_data)
         output

     ():
        
         inference_output.cpu().numpy().tolist()

     ():
        
        :
            data = .preprocess(data)
            data = data.to(.device)
            output = .inference(data)
             .postprocess(output)
         Exception  e:
             [{: (e)}]

Archive and Serve:

# Archive model
torch-model-archiver \
  --model-name mymodel \
  --version 1.0 \
  --serialized-file model.pth \
  --handler handler.py \
  --extra-files config.json,index_to_name.json \
  --export-path model_store

# Start TorchServe
torchserve --start --ncs --model-store model_store --models mymodel=mymodel.mar

# Make prediction
curl -X POST http://localhost:8080/predictions/mymodel \
  -H "Content-Type: application/json" \
  -d '{"data": [[...]]}'

Model Optimization

Quantization

Quantization reduces model size and improves inference speed by using lower precision numbers.

Post-Training Quantization (PTQ):

import torch
from torch.quantization import quantize_dynamic

# Dynamic quantization (weights quantized, activations computed in float)
model = MyModel()

# Quantize specific layers
quantized_model = quantize_dynamic(
    model,
    {nn.Linear, nn.LSTM},  # Layers to quantize
    dtype=torch.qint8      # Quantization dtype
)

# Save quantized model
torch.jit.save(torch.jit.script(quantized_model), "model_quantized.pt")

Static Quantization:

import torch
from torch.quantization import (
    quantize,
    prepare,
    convert,
    get_default_qconfig,
)

# Prepare model for static quantization
model = MyModel()
model.eval()

# Set quantization configuration
model.qconfig = get_default_qconfig('fbgemm')

# Prepare model with calibration data
prepared_model = prepare(model)

# Calibrate with representative data
with torch.no_grad():
    for data in calibration_dataloader:
        prepared_model(data)

# Convert to quantized model
quantized_model = convert(prepared_model)

# Save
torch.jit.save(torch.jit.script(quantized_model), "model_static_quantized.pt")

Quantization-Aware Training (QAT):

import torch
from torch.quantization import prepare_qat, convert

# Prepare model for QAT
model = MyModel()
model.train()
model.qconfig = get_default_qconfig('fbgemm')

# Prepare for QAT
model_prepared = prepare_qat(model, inplace=True)

# Fine-tune with quantization simulation
optimizer = torch.optim.SGD(model_prepared.parameters(), lr=0.01)

for epoch in range(num_epochs):
    for batch in train_dataloader:
        optimizer.zero_grad()
        loss = criterion(model_prepared(batch[0]), batch[1])
        loss.backward()
        optimizer.step()

# Convert to quantized model
model_prepared.eval()
quantized_model = convert(model_prepared)

Pruning

Pruning removes less important weights from model.

Structured Pruning:

import torch.nn.utils.prune as prune
import torch

model = MyModel()

# Prune 30% of weights in linear layers
for name, module in model.named_modules():
    if isinstance(module, nn.Linear):
        prune.l1_unstructured(module, name='weight', amount=0.3)

# Make pruning permanent
for name, module in model.named_modules():
    if isinstance(module, nn.Linear):
        prune.remove(module, 'weight')

Global Unstructured Pruning:

# Prune across all layers globally
parameters_to_prune = []
for name, module in model.named_modules():
    if isinstance(module, (nn.Linear, nn.Conv2d)):
        parameters_to_prune.append((module, 'weight'))

prune.global_unstructured(
    parameters_to_prune,
    pruning_method=prune.L1Unstructured,
    amount=0.2
)

Iterative Pruning:

def iterative_pruning(model, train_loader, num_iterations=5, prune_amount=0.2):
    for iteration in range(num_iterations):
        print(f"Pruning iteration {iteration + 1}/{num_iterations}")

        # Prune
        for name, module in model.named_modules():
            if isinstance(module, nn.Linear):
                prune.l1_unstructured(module, name='weight', amount=prune_amount)

        # Fine-tune
        for epoch in range(5):
            for batch in train_loader:
                # Training code here
                pass

    # Make pruning permanent
    for name, module in model.named_modules():
        if isinstance(module, nn.Linear):
            prune.remove(module, 'weight')

    return model

Model Compression

Knowledge Distillation:

import torch
import torch.nn as nn
import torch.nn.functional as F

class DistillationLoss(nn.Module):
    def __init__(self, alpha=0.5, temperature=2.0):
        super().__init__()
        self.alpha = alpha
        self.temperature = temperature
        self.kl_div = nn.KLDivLoss(reduction='batchmean')

    def forward(self, student_logits, teacher_logits, targets):
        # Hard loss (cross-entropy with true labels)
        hard_loss = F.cross_entropy(student_logits, targets)

        # Soft loss (KL divergence with teacher)
        soft_loss = self.kl_div(
            F.log_softmax(student_logits / self.temperature, dim=1),
            F.softmax(teacher_logits / self.temperature, dim=1)
        ) * (self.temperature ** 2)

        return self.alpha * soft_loss + (1 - self.alpha) * hard_loss

# Training loop
teacher_model = load_teacher_model()
student_model = create_student_model()
criterion = DistillationLoss(alpha=0.7, temperature=3.0)

teacher_model.eval()
student_model.train()

for batch in train_loader:
    inputs, targets = batch

    with torch.no_grad():
        teacher_outputs = teacher_model(inputs)

    student_outputs = student_model(inputs)
    loss = criterion(student_outputs, teacher_outputs, targets)

    loss.backward()
    optimizer.step()

Inference Optimization

Batching

import torch
from collections import deque
import threading
import time

class BatchInferenceServer:
    def __init__(self, model, max_batch_size=32, max_wait_time=0.1):
        self.model = model
        self.model.eval()
        self.max_batch_size = max_batch_size
        self.max_wait_time = max_wait_time
        self.batch_queue = deque()
        self.results = {}
        self.lock = threading.Lock()
        self.running = False

    def start(self):
        self.running = True
        self.thread = threading.Thread(target=self._process_batches)
        self.thread.start()

    def stop(self):
        self.running = False
        self.thread.join()

    def predict(self, input_data):
        request_id = id(input_data)
        with self.lock:
            self.batch_queue.append((request_id, input_data))
        return request_id

     ():
        start_time = time.time()
         request_id   .results:
             time.time() - start_time > timeout:
                 TimeoutError()
            time.sleep()
         .results.pop(request_id)

     ():
         .running:
            batch = []
            start_time = time.time()

             .lock:
                 (batch) < .max_batch_size  \
                      (time.time() - start_time) < .max_wait_time:
                     .batch_queue:
                        batch.append(.batch_queue.popleft())
                    :
                        time.sleep()

             batch:
                request_ids, inputs = (*batch)
                batch_tensor = torch.stack(inputs)

                 torch.no_grad():
                    outputs = .model(batch_tensor)

                 .lock:
                     req_id, output  (request_ids, outputs):
                        .results[req_id] = output

GPU Utilization

import torch
import torch.multiprocessing as mp

def run_inference(rank, model, inputs, outputs):
    """Worker function for multi-GPU inference."""
    torch.cuda.set_device(rank)
    model = model.to(rank)
    model.eval()

    with torch.no_grad():
        outputs[rank] = model(inputs[rank])

def multi_gpu_inference(model, inputs):
    """Distribute inference across multiple GPUs."""
    num_gpus = torch.cuda.device_count()
    outputs = [None] * num_gpus

    # Split inputs across GPUs
    inputs_per_gpu = torch.chunk(inputs, num_gpus)
    inputs = [inp.to(i) for i, inp in enumerate(inputs_per_gpu)]

    # Spawn processes
    mp.spawn(
        run_inference,
        args=(model, inputs, outputs),
        nprocs=num_gpus,
        join=True
    )

    return torch.cat(outputs, dim=0)

Mixed Precision

import torch
from torch.cuda.amp import autocast, GradScaler

# For inference with mixed precision
def mixed_precision_inference(model, inputs):
    model.eval()
    with autocast():
        with torch.no_grad():
            outputs = model(inputs)
    return outputs

# For training with mixed precision
scaler = GradScaler()

for batch in train_loader:
    inputs, targets = batch
    inputs, targets = inputs.to(device), targets.to(device)

    optimizer.zero_grad()

    with autocast():
        outputs = model(inputs)
        loss = criterion(outputs, targets)

    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

Deployment Patterns

FastAPI Server

from fastapi import FastAPI, File, UploadFile, HTTPException
from fastapi.responses import JSONResponse
from pydantic import BaseModel
import torch
import torch.nn as nn
from torchvision import transforms
from PIL import Image
import io

app = FastAPI(title="PyTorch Model API")

# Load model
class ImageClassifier(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 64, 3)
        self.pool = nn.MaxPool2d(2, 2)
        self.fc = nn.Linear(64 * 13 * 13, 10)

    def forward(self, x):
        x = self.pool(torch.relu(self.conv1(x)))
        x = x.view(x.size(0), -1)
        return self.fc(x)

model = ImageClassifier()
model.load_state_dict(torch.load("model.pth"))
model.eval()

# Preprocessing
transform = transforms.Compose([
    transforms.Resize((28, 28)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[, , ], std=[, , ])
])

 ():
    class_id: 
    class_name: 
    confidence: 


  ():
    :
        
        image_data =  file.read()
        image = Image.(io.BytesIO(image_data)).convert()
        input_tensor = transform(image).unsqueeze()

        
         torch.no_grad():
            outputs = model(input_tensor)
            probabilities = torch.softmax(outputs, dim=)
            confidence, predicted = torch.(probabilities, )

         PredictionResponse(
            class_id=predicted.item(),
            class_name=,
            confidence=confidence.item()
        )
     Exception  e:
         HTTPException(status_code=, detail=(e))


  ():
     {: }


  ():
     {
        : ,
        : (p.numel()  p  model.parameters()),
        : ,
        : 
    }

 __name__ == :
     uvicorn
    uvicorn.run(app, host=, port=)

TorchServe Configuration

config.properties:

inference_address=http://0.0.0.0:8080
management_address=http://0.0.0.0:8081
metrics_address=http://0.0.0.0:8082
number_of_netty_threads=4
job_queue_size=10
model_store=model_store
load_models=all
number_of_gpu=1
default_response_timeout=120

Docker Deployment:

FROM pytorch/torchserve:latest

# Copy model archive
COPY model_store /home/model-server/model-store

# Copy config
COPY config.properties /home/model-server/config.properties

# Expose ports
EXPOSE 8080 8081 8082

# Start TorchServe
CMD ["torchserve", \
     "--start", \
     "--model-store", "/home/model-server/model-store", \
     "--models", "mymodel=mymodel.mar", \
     "--ts-config", "/home/model-server/config.properties"]

ONNX Runtime Server

Python ONNX Runtime Server:

from fastapi import FastAPI
import numpy as np
import onnxruntime as ort
from pydantic import BaseModel

app = FastAPI()

# Load ONNX model
session = ort.InferenceSession("model.onnx")

class InputData(BaseModel):
    data: list

@app.post("/predict")
async def predict(input_data: InputData):
    input_array = np.array(input_data.data, dtype=np.float32)

    # Run inference
    outputs = session.run(
        None,
        {session.get_inputs()[0].name: input_array}
    )

    return {"output": outputs[0].tolist()}

Model Versioning

Versioning Strategy

import os
import json
from datetime import datetime
import torch

class ModelVersionManager:
    def __init__(self, base_path="models"):
        self.base_path = base_path
        os.makedirs(base_path, exist_ok=True)

    def save_model(self, model, version, metadata=None):
        """Save model with version and metadata."""
        version_path = os.path.join(self.base_path, f"v{version}")
        os.makedirs(version_path, exist_ok=True)

        # Save model weights
        model_path = os.path.join(version_path, "model.pth")
        torch.save(model.state_dict(), model_path)

        # Save metadata
        metadata = metadata or {}
        metadata.update({
            "version": version,
            "saved_at": datetime.now().isoformat(),
            "model_path": model_path
        })

        metadata_path = os.path.join(version_path, "metadata.json")
        with open(metadata_path, "w") as f:
            json.dump(metadata, f, indent=2)

        return version_path

    def load_model(self, version, model_class):
        """Load model by version."""
        version_path = os.path.join(self.base_path, f"v")
        model_path = os.path.join(version_path, )

        model = model_class()
        model.load_state_dict(torch.load(model_path))
        model.()

         model

     ():
        
        versions = []
         item  os.listdir(.base_path):
             item.startswith():
                version_path = os.path.join(.base_path, item)
                metadata_path = os.path.join(version_path, )
                 os.path.exists(metadata_path):
                     (metadata_path)  f:
                        versions.append(json.load(f))
         (versions, key= x: x[])

A/B Testing Models

import random
from typing import Dict, Optional
import torch

class ABTestModelRouter:
    def __init__(self, models: Dict[str, torch.nn.Module], traffic_split: Dict[str, float]):
        """
        Args:
            models: Dictionary of model_name -> model
            traffic_split: Dictionary of model_name -> traffic_percentage (sum must be 1.0)
        """
        self.models = models
        self.traffic_split = traffic_split
        self.model_names = list(traffic_split.keys())
        self.cumulative_split = []
        cumulative = 0
        for name in self.model_names:
            cumulative += traffic_split[name]
            self.cumulative_split.append(cumulative)

    def get_model(self, request_id: Optional[str] = None) -> torch.nn.Module:
        """Select model based on traffic split."""
        # Use request_id for consistent routing (same request always goes to same model)
        if request_id:
            hash_val = hash(request_id) % 1000
            rand_val = hash_val / 1000.0
        else:
            rand_val = random.random()

        for i, threshold in (.cumulative_split):
             rand_val < threshold:
                 .models[.model_names[i]]

         .models[.model_names[-]]

     ():
        
        model = .get_model(request_id)
        model.()

         torch.no_grad():
            output = model(input_data)

         output


model_a = create_model_v1()
model_b = create_model_v2()

router = ABTestModelRouter(
    models={: model_a, : model_b},
    traffic_split={: , : }
)


output = router.predict(input_data, request_id=)

Model Monitoring

import time
import json
from collections import defaultdict
from datetime import datetime
import torch

class ModelMonitor:
    def __init__(self, model_name: str):
        self.model_name = model_name
        self.metrics = defaultdict(list)
        self.start_time = time.time()

    def log_prediction(self, request_id: str, input_shape: tuple,
                       output_shape: tuple, latency: float,
                       model_version: str):
        """Log prediction metrics."""
        self.metrics["predictions"].append({
            "request_id": request_id,
            "timestamp": datetime.now().isoformat(),
            "input_shape": input_shape,
            "output_shape": output_shape,
            "latency_ms": latency,
            "model_version": model_version
        })

    def log_error(self, request_id: str, error_type: str, error_message: str):
        """Log prediction errors."""
        self.metrics["errors"].append({
            "request_id": request_id,
            : datetime.now().isoformat(),
            : error_type,
            : error_message
        })

     ():
        
        predictions = .metrics[]
        errors = .metrics[]

         predictions:
            avg_latency = (p[]  p  predictions) / (predictions)
            total_predictions = (predictions)
        :
            avg_latency = 
            total_predictions = 

         {
            : .model_name,
            : time.time() - .start_time,
            : total_predictions,
            : (errors),
            : avg_latency,
            : (errors) / (total_predictions, ) * 
        }


monitor = ModelMonitor()


  ():
    request_id = (uuid.uuid4())
    start_time = time.time()

    :
        
        output = model(input_tensor)
        latency = (time.time() - start_time) * 

        monitor.log_prediction(
            request_id=request_id,
            input_shape=(input_tensor.shape),
            output_shape=(output.shape),
            latency=latency,
            model_version=
        )

         {: output.tolist()}

     Exception  e:
        monitor.log_error(request_id, (e).__name__, (e))
         HTTPException(status_code=, detail=(e))

Error Handling

import logging
from functools import wraps
import torch

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class ModelInferenceError(Exception):
    """Base exception for model inference errors."""
    pass

class ModelLoadError(ModelInferenceError):
    """Exception raised when model fails to load."""
    pass

class InputValidationError(ModelInferenceError):
    """Exception raised when input validation fails."""
    pass

def handle_inference_errors(func):
    """Decorator for handling inference errors."""
    @wraps(func)
    def wrapper(*args, **kwargs):
        try:
            return func(*args, **kwargs)
        except ModelLoadError as e:
            logger.error(f"Model load error: {e}")
            raise
        except InputValidationError as e:
            logger.warning(f"Input validation error: {e}")
            raise
        except torch.cuda.OutOfMemoryError:
            logger.error("CUDA out of memory")
            raise ModelInferenceError("GPU memory exhausted")
         Exception  e:
            logger.error()
             ModelInferenceError()
     wrapper

 :
    
     ():
        .model_path = model_path
        .device = device
        .model = 
        ._load_model()

     ():
        
        :
            .model = torch.load(.model_path, map_location=.device)
            .model.()
            logger.info()
         FileNotFoundError:
             ModelLoadError()
         Exception  e:
             ModelLoadError()


     ():
        
          (input_data, torch.Tensor):
             InputValidationError()

         input_data.dim() != :
             InputValidationError()

         torch.no_grad():
            output = .model(input_data.to(.device))

         output.cpu()

     ():
        
        total_params = (p.numel()  p  .model.parameters())
        trainable_params = (p.numel()  p  .model.parameters()  p.requires_grad)

         {
            : .model_path,
            : (.device),
            : total_params,
            : trainable_params
        }

Performance Benchmarking

import time
import torch
import numpy as np
from typing import List, Dict
import json

class ModelBenchmark:
    def __init__(self, model, input_shape, warmup_runs=10, benchmark_runs=100):
        self.model = model
        self.model.eval()
        self.input_shape = input_shape
        self.warmup_runs = warmup_runs
        self.benchmark_runs = benchmark_runs
        self.device = next(model.parameters()).device

    def _generate_input(self, batch_size=1):
        """Generate random input for benchmarking."""
        return torch.randn(batch_size, *self.input_shape, device=self.device)

    def benchmark_latency(self, batch_sizes: List[int] = [1, 8, 16, 32]):
        """Benchmark inference latency for different batch sizes."""
        results = {}

        for batch_size in batch_sizes:
            # Warmup
            for _ in (.warmup_runs):
                input_data = ._generate_input(batch_size)
                 torch.no_grad():
                    _ = .model(input_data)

            
            latencies = []
             _  (.benchmark_runs):
                input_data = ._generate_input(batch_size)

                torch.cuda.synchronize()  .device. ==   
                start_time = time.perf_counter()

                 torch.no_grad():
                    _ = .model(input_data)

                torch.cuda.synchronize()  .device. ==   
                end_time = time.perf_counter()

                latencies.append((end_time - start_time) * )  

            results[batch_size] = {
                : np.mean(latencies),
                : np.std(latencies),
                : np.(latencies),
                : np.(latencies),
                : np.percentile(latencies, ),
                : np.percentile(latencies, ),
                : np.percentile(latencies, )
            }

         results

     ():
        
        input_data = ._generate_input()

        start_time = time.time()
        predictions = 

         (time.time() - start_time) < duration_seconds:
             torch.no_grad():
                _ = .model(input_data)
            predictions += 

        elapsed = time.time() - start_time
        throughput = predictions / elapsed

         {
            : elapsed,
            : predictions,
            : throughput
        }

     ():
        
         .device. != :
             {: }

        results = {}
        torch.cuda.reset_peak_memory_stats()

         batch_size  batch_sizes:
            torch.cuda.empty_cache()
            torch.cuda.reset_peak_memory_stats()

            input_data = ._generate_input(batch_size)

             torch.no_grad():
                _ = .model(input_data)

            results[batch_size] = {
                : torch.cuda.max_memory_allocated() /  / ,
                : torch.cuda.max_memory_reserved() /  / 
            }

         results

     ():
        
        ()
        ( * )

        
        ()
        latency_results = .benchmark_latency()
         batch_size, metrics  latency_results.items():
            ()

        
        ()
        throughput_results = .benchmark_throughput()
        ()

        
        ()
        memory_results = .benchmark_memory()
         batch_size, metrics  memory_results.items():
            ()

        
        ()
        total_params = (p.numel()  p  .model.parameters())
        ()

         {
            : latency_results,
            : throughput_results,
            : memory_results,
            : total_params
        }


model = load_model()
benchmark = ModelBenchmark(model, input_shape=(, , ))
results = benchmark.run_full_benchmark()


 (, )  f:
    json.dump(results, f, indent=)

Best Practices

Pre-Deployment Checklist

• Model Export

  • Model exported to production format (TorchScript/ONNX)
  • Exported model tested and verified
  • Model size optimized (quantization/pruning)

• Performance

  • Inference latency meets SLA (< 100ms for real-time)
  • Throughput tested with expected load
  • GPU memory usage optimized
  • Batch processing configured

• Reliability

  • Error handling implemented
  • Graceful degradation for failures
  • Circuit breaker pattern for external dependencies
  • Retry logic for transient failures

• Monitoring

  • Metrics collection (latency, throughput, errors)
  • Logging configured
  • Health check endpoint
  • Alert thresholds set

• Security

  • Input validation implemented
  • Rate limiting configured
  • Authentication/authorization for API
  • Model files stored securely

• Deployment

  • Docker container created
  • Environment variables configured
  • CI/CD pipeline set up
  • Blue-green deployment strategy

Post-Deployment Checklist

• Validation

  • Smoke tests passed
  • A/B test started
  • Model performance monitored
  • Error rates within acceptable range

• Documentation

  • API documentation updated
  • Model version documented
  • Known issues documented
  • Runbook created

Performance Optimization Tips

1. Use TorchScript for Production

   • Export models to TorchScript for faster inference
   • Use tracing for models without control flow
   • Use scripting for models with dynamic control flow

2. Apply Quantization

   • Use dynamic quantization for quick deployment
   • Use static quantization for better performance
   • Use QAT for minimal accuracy loss

3. Optimize Batch Size

   • Find optimal batch size for your hardware
   • Use larger batches for better GPU utilization
   • Consider latency requirements when choosing batch size

4. Use Mixed Precision

   • Enable FP16 for faster computation
   • Use GradScaler for training stability
   • Test accuracy impact before deployment

5. Monitor Model Performance

   • Track latency, throughput, and error rates
   • Set up alerts for performance degradation
   • Monitor GPU memory usage
   • Track prediction drift

Related Skills

• 05-ai-ml-core/model-training
• 05-ai-ml-core/model-optimization
• 05-ai-ml-core/model-versioning
• 06-ai-ml-production/llm-integration
• 15-devops-infrastructure/ci-cd-pipelines