1 Why REST APIs for Model Serving

Once a model is trained and validated, it needs to be consumed by something: a web app's checkout flow, a mobile app's recommendation feed, a batch job scoring a million rows overnight, or another internal microservice. Each of these consumers might be written in a different language, deployed on different infrastructure, and updated on a different schedule than your model. You need a contract that is language-agnostic, network-accessible, and independently versionable. REST over HTTP is the dominant choice because almost every platform on earth can make an HTTP request.

The Alternatives, and Why REST Usually Wins

• Embed the model directly in the consuming app — works for a single Python app, but forces every consumer to install your exact ML stack (PyTorch, scikit-learn, CUDA drivers) and forces you to redeploy every consumer whenever the model changes. Doesn't scale past one team.
• Batch scoring to a database/file — great for offline use cases (nightly churn scores), but useless when you need a prediction synchronously, e.g., "should this transaction be approved right now?"
• gRPC — lower latency and strongly-typed contracts (via Protocol Buffers), genuinely better for very high-throughput internal service-to-service calls. The tradeoff: harder to debug by hand, no built-in browser tooling, and a steeper learning curve. Many companies start with REST and migrate hot paths to gRPC later.
• Message queues (Kafka, SQS) — excellent for asynchronous, decoupled, high-volume event-driven scoring, but adds operational complexity and isn't a fit when a caller needs an immediate response.

REST's winning combination for most ML serving: human-readable (JSON), debuggable with curl in ten seconds, supported natively by literally every programming language and every cloud load balancer, and easy to document and version. It is the right default; reach for gRPC or queues only when you've measured a concrete need (sub-millisecond latency budgets, very high request rates, or a genuinely asynchronous workflow).

What "Serving" Actually Involves

Model serving is more than calling model.predict(). A production serving layer must: validate that the incoming request actually contains the features the model expects, in the right types and ranges; load the model once and keep it warm in memory rather than reloading it on every request; handle concurrent requests without one slow prediction blocking all others; report whether it's healthy so a load balancer or Kubernetes can route traffic away from a broken instance; and fail predictably and informatively when something goes wrong, rather than crashing the process. FastAPI plus Docker gives you a clean, minimal way to get all of this right.

2 Designing the Prediction API: Schemas and Versioning

Before writing a single route, design the contract. A good ML API contract specifies exactly what fields a request must contain, their types and valid ranges, what the response looks like on success, and what it looks like on failure. Pydantic models are how FastAPI expresses this contract in code — and that same code doubles as runtime validation and as auto-generated documentation.

We'll build a churn-prediction API for a subscription business: given a customer's usage features, predict the probability they cancel next month. The model itself was trained in earlier lessons (a gradient-boosted classifier) and saved with joblib, or registered in MLflow as discussed in Lesson 65.

Request and Response Schemas

# app/schemas.py
from pydantic import BaseModel, Field, field_validator
from typing import Literal


class ChurnFeatures(BaseModel):
    """Input schema for a single churn prediction request."""
    tenure_months: int = Field(..., ge=0, le=600, description="Months as a customer")
    monthly_charges: float = Field(..., ge=0, description="Current monthly bill in USD")
    total_charges: float = Field(..., ge=0, description="Lifetime billed amount in USD")
    contract_type: Literal["month-to-month", "one-year", "two-year"]
    num_support_tickets: int = Field(..., ge=0, le=200)
    has_autopay: bool

    @field_validator("total_charges")
    @classmethod
    def total_at_least_monthly(cls, v, info):
        # Cross-field sanity check: total can't be less than one month's charge
        monthly = info.data.get("monthly_charges")
        if monthly is not None and v < monthly:
            raise ValueError("total_charges cannot be less than monthly_charges")
        return v

    model_config = {
        "json_schema_extra": {
            "example": {
                "tenure_months": 14,
                "monthly_charges": 79.99,
                "total_charges": 1119.86,
                "contract_type": "month-to-month",
                "num_support_tickets": 3,
                "has_autopay": False,
            }
        }
    }


class PredictionResponse(BaseModel):
    """Output schema for a single prediction."""
    churn_probability: float = Field(..., ge=0.0, le=1.0)
    will_churn: bool
    model_version: str
    risk_tier: Literal["low", "medium", "high"]


class BatchPredictionRequest(BaseModel):
    """Input schema for scoring many customers in one request."""
    customers: list[ChurnFeatures] = Field(..., min_length=1, max_length=1000)


class BatchPredictionResponse(BaseModel):
    predictions: list[PredictionResponse]
    count: int


class ErrorResponse(BaseModel):
    error: str
    detail: str

Notice what this buys you for free: tenure_months can never be negative, contract_type can never be a typo'd string outside the three valid values, and the cross-field validator catches a logically inconsistent request before it ever reaches the model. None of this required writing an if statement in the route handler — Pydantic raises a 422 Unprocessable Entity automatically with a precise message about which field failed and why.

Why Version the Endpoint Path

Putting /v1/ in front of /predict looks like overkill on day one, but it is the cheapest insurance you can buy. When you retrain the model with a different feature set six months from now — say you add a support_sentiment_score feature — you can stand up /v2/predict with the new contract while /v1/predict keeps serving the old contract for clients that haven't migrated yet. Without versioning, every model or schema change is a breaking change for every consumer simultaneously.

3 Building the FastAPI Service

The service has three jobs: load the model once at startup, expose a /v1/predict route that validates input and returns a structured prediction, and expose a /health route for orchestrators. We'll support loading the model either from a local joblib file or by pulling the current "Production"-stage model from the MLflow Model Registry you set up in Lesson 65 — selectable via an environment variable, so the same image works in a laptop demo and a real deployment.

# app/model_loader.py
import os
import joblib
import mlflow
import mlflow.sklearn

MODEL_SOURCE = os.getenv("MODEL_SOURCE", "local")          # "local" or "mlflow"
LOCAL_MODEL_PATH = os.getenv("LOCAL_MODEL_PATH", "models/churn_model.joblib")
MLFLOW_TRACKING_URI = os.getenv("MLFLOW_TRACKING_URI", "http://mlflow:5000")
MLFLOW_MODEL_NAME = os.getenv("MLFLOW_MODEL_NAME", "churn-classifier")
MLFLOW_MODEL_STAGE = os.getenv("MLFLOW_MODEL_STAGE", "Production")


def load_model():
    """
    Load the serving model from local disk or from the MLflow Model Registry,
    depending on configuration. Returns (model, model_version_string).
    """
    if MODEL_SOURCE == "mlflow":
        mlflow.set_tracking_uri(MLFLOW_TRACKING_URI)
        model_uri = f"models:/{MLFLOW_MODEL_NAME}/{MLFLOW_MODEL_STAGE}"
        model = mlflow.sklearn.load_model(model_uri)

        # Resolve the concrete version number behind the "Production" alias
        # so the API can report exactly which model answered each request.
        client = mlflow.tracking.MlflowClient()
        version_info = client.get_latest_versions(MLFLOW_MODEL_NAME, stages=[MLFLOW_MODEL_STAGE])
        version = version_info[0].version if version_info else "unknown"
        return model, f"mlflow:{MLFLOW_MODEL_NAME}:v{version}"

    # Default: load a joblib-serialized model from local disk
    model = joblib.load(LOCAL_MODEL_PATH)
    return model, f"local:{os.path.basename(LOCAL_MODEL_PATH)}"

# app/main.py
import logging
import time
from contextlib import asynccontextmanager

import numpy as np
import pandas as pd
from fastapi import FastAPI, HTTPException, Request, status
from fastapi.responses import JSONResponse

from app.model_loader import load_model
from app.schemas import (
    ChurnFeatures, PredictionResponse,
    BatchPredictionRequest, BatchPredictionResponse,
)

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("churn-api")

ml_state = {}   # holds the loaded model + version, set once at startup


@asynccontextmanager
async def lifespan(app: FastAPI):
    # Startup: load the model exactly once, before accepting traffic
    logger.info("Loading model...")
    model, version = load_model()
    ml_state["model"] = model
    ml_state["version"] = version
    logger.info(f"Model loaded: {version}")
    yield
    # Shutdown: nothing to clean up for a stateless sklearn model
    ml_state.clear()


app = FastAPI(
    title="Churn Prediction API",
    version="1.0.0",
    lifespan=lifespan,
)


def features_to_dataframe(features: ChurnFeatures) -> pd.DataFrame:
    """Convert a validated Pydantic object into the row shape the model expects."""
    row = features.model_dump()
    return pd.DataFrame([row])


def risk_tier_for(probability: float) -> str:
    if probability < 0.3:
        return "low"
    if probability < 0.7:
        return "medium"
    return "high"


def run_prediction(features: ChurnFeatures) -> PredictionResponse:
    model = ml_state["model"]
    df = features_to_dataframe(features)
    probability = float(model.predict_proba(df)[0, 1])
    return PredictionResponse(
        churn_probability=round(probability, 4),
        will_churn=probability >= 0.5,
        model_version=ml_state["version"],
        risk_tier=risk_tier_for(probability),
    )


@app.get("/health")
def health():
    """Lightweight liveness/readiness check for orchestrators and load balancers."""
    return {"status": "ok", "model_loaded": "model" in ml_state, "model_version": ml_state.get("version")}


@app.post("/v1/predict", response_model=PredictionResponse)
def predict(features: ChurnFeatures):
    """Score a single customer. Synchronous: the model call is CPU-bound and fast (<10ms)."""
    try:
        return run_prediction(features)
    except Exception as exc:
        logger.exception("Prediction failed")
        raise HTTPException(status_code=500, detail="Model inference failed") from exc


@app.post("/v1/predict/batch", response_model=BatchPredictionResponse)
def predict_batch(request: BatchPredictionRequest):
    """Score up to 1000 customers in a single request — one model call, not N."""
    try:
        model = ml_state["model"]
        df = pd.DataFrame([c.model_dump() for c in request.customers])
        probabilities = model.predict_proba(df)[:, 1]

        predictions = [
            PredictionResponse(
                churn_probability=round(float(p), 4),
                will_churn=p >= 0.5,
                model_version=ml_state["version"],
                risk_tier=risk_tier_for(float(p)),
            )
            for p in probabilities
        ]
        return BatchPredictionResponse(predictions=predictions, count=len(predictions))
    except Exception as exc:
        logger.exception("Batch prediction failed")
        raise HTTPException(status_code=500, detail="Batch inference failed") from exc

Run this locally with uvicorn app.main:app --reload --port 8000 and visit http://localhost:8000/docs — FastAPI auto-generates an interactive Swagger UI directly from your Pydantic schemas, so anyone on your team can try the API without writing a line of client code.

4 Async vs Sync Endpoints, and Error Handling

When async Actually Helps

FastAPI lets you declare a route as async def or plain def. The common misconception is that async automatically makes inference faster — it does not. A scikit-learn or PyTorch model.predict() call is CPU-bound: it occupies the Python interpreter doing matrix math, and awaiting it doesn't free up the event loop because there's nothing to await. For purely CPU-bound inference, a plain synchronous def predict(...) route is correct — FastAPI automatically runs sync routes in a thread pool, so one slow prediction doesn't block the event loop from accepting new connections.

async def earns its keep when the route does genuine I/O: calling an external feature store over the network, querying a database for a customer's history before scoring, or calling another microservice. In those cases, awaiting the I/O lets the event loop serve other requests while waiting, which a blocking sync call cannot do.

# Sync route — correct for pure CPU-bound model inference
@app.post("/v1/predict")
def predict(features: ChurnFeatures):
    return run_prediction(features)   # FastAPI runs this in a thread pool automatically


# Async route — correct when you must await I/O before/after inference
@app.post("/v1/predict/enriched")
async def predict_enriched(customer_id: str, features: ChurnFeatures):
    # Example: fetch additional features from an external store before scoring
    history = await feature_store_client.get_customer_history(customer_id)
    enriched = features.model_dump() | {"avg_session_minutes": history["avg_session_minutes"]}
    probability = float(ml_state["model"].predict_proba(pd.DataFrame([enriched]))[0, 1])
    return {"churn_probability": round(probability, 4)}

Structured Error Handling

FastAPI returns a 422 Unprocessable Entity automatically whenever a request fails Pydantic validation — you don't write that path yourself. Your job is the rest: turning unexpected exceptions during inference into clean 500 responses instead of stack traces leaking to the client, and giving callers a consistent JSON error shape they can branch on programmatically.

# app/main.py (continued) — global exception handlers
from fastapi.exceptions import RequestValidationError

@app.exception_handler(RequestValidationError)
async def validation_exception_handler(request: Request, exc: RequestValidationError):
    # Reshape FastAPI's default 422 body into our own ErrorResponse contract
    first_error = exc.errors()[0]
    field = ".".join(str(loc) for loc in first_error["loc"] if loc != "body")
    return JSONResponse(
        status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
        content={"error": "validation_error", "detail": f"{field}: {first_error['msg']}"},
    )


@app.exception_handler(Exception)
async def unhandled_exception_handler(request: Request, exc: Exception):
    logger.exception("Unhandled exception")
    return JSONResponse(
        status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
        content={"error": "internal_error", "detail": "An unexpected error occurred. Please retry or contact support."},
    )

5 Health Checks and API Key Authentication

Health Checks for Orchestrators

Kubernetes, AWS ECS, and basically every container orchestrator periodically probes a known endpoint to decide whether an instance is healthy enough to receive traffic (a readiness probe) and whether it's alive enough to keep running at all (a liveness probe). The /health route from Section 3 already serves this purpose; the key design point is that it must be fast (no model inference, no expensive database query) and it must accurately reflect whether the service can actually serve a prediction right now — for example, returning a 503 if the model failed to load at startup rather than always returning 200.

# app/main.py (revised health check)
from fastapi import Response

@app.get("/health")
def health(response: Response):
    model_ready = "model" in ml_state
    if not model_ready:
        response.status_code = status.HTTP_503_SERVICE_UNAVAILABLE
    return {
        "status": "ok" if model_ready else "unavailable",
        "model_loaded": model_ready,
        "model_version": ml_state.get("version"),
    }

Protecting the API with an API Key

Full OAuth2/JWT flows are usually overkill for an internal service-to-service prediction API. A simple, effective pattern is a shared API key sent in a header, checked via a FastAPI dependency. The dependency is reusable across every protected route and keeps auth logic out of your business code.

# app/auth.py
import os
from fastapi import Security, HTTPException, status
from fastapi.security import APIKeyHeader

API_KEY = os.environ["API_KEY"]               # fail fast if not configured
api_key_header = APIKeyHeader(name="X-API-Key", auto_error=False)


def require_api_key(provided_key: str = Security(api_key_header)):
    if provided_key != API_KEY:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Missing or invalid API key",
        )

# app/main.py — apply the dependency to protected routes only
from fastapi import Depends
from app.auth import require_api_key

@app.post("/v1/predict", response_model=PredictionResponse, dependencies=[Depends(require_api_key)])
def predict(features: ChurnFeatures):
    return run_prediction(features)

# /health is deliberately left unprotected — orchestrators probing it
# shouldn't need to manage a secret, and it leaks no sensitive data.

6 Dockerizing the Service: Multi-Stage Builds

Docker packages the API, its exact Python dependencies, and a minimal OS into a single portable image. The same image that passes your tests on a laptop is the image that runs in production — eliminating an entire class of "works on my machine" bugs. A naive Dockerfile works, but a multi-stage build produces a meaningfully smaller, more secure final image by separating "things needed to install dependencies" (compilers, build tools) from "things needed to run the app" (just the installed packages and your code).

Pinning Dependencies

Before writing the Dockerfile, pin exact versions in requirements.txt. Unpinned dependencies (fastapi instead of fastapi==0.111.0) mean a routine rebuild six months from now can silently pull a newer, incompatible version of a library and break the service in production with zero code changes on your part.

# requirements.txt
fastapi==0.111.0
uvicorn[standard]==0.30.1
pydantic==2.7.4
scikit-learn==1.5.0
pandas==2.2.2
numpy==1.26.4
joblib==1.4.2
mlflow==2.14.1

The Multi-Stage Dockerfile

# Dockerfile
# ---- Stage 1: builder — installs dependencies into a virtualenv ----
FROM python:3.11-slim AS builder

WORKDIR /build

# Build tools needed only to compile some wheels (e.g. scientific libs);
# these never end up in the final image.
RUN apt-get update && apt-get install -y --no-install-recommends \
        build-essential \
    && rm -rf /var/lib/apt/lists/*

COPY requirements.txt .
RUN python -m venv /opt/venv
ENV PATH="/opt/venv/bin:$PATH"
RUN pip install --no-cache-dir --upgrade pip \
    && pip install --no-cache-dir -r requirements.txt


# ---- Stage 2: runtime — slim image with only what's needed to run ----
FROM python:3.11-slim AS runtime

# Run as a non-root user — a basic but important container security practice
RUN useradd --create-home --uid 1000 appuser

WORKDIR /app

# Copy the pre-built virtualenv from the builder stage (no compilers here)
COPY --from=builder /opt/venv /opt/venv
ENV PATH="/opt/venv/bin:$PATH"

COPY app/ ./app/
COPY models/ ./models/

USER appuser

EXPOSE 8000

HEALTHCHECK --interval=30s --timeout=3s --start-period=10s --retries=3 \
    CMD python -c "import urllib.request; urllib.request.urlopen('http://localhost:8000/health')" || exit 1

CMD ["uvicorn", "app.main:app", "--host", "0.0.0.0", "--port", "8000", "--workers", "2"]

The builder stage installs build-essential (gcc, headers — needed to compile some scientific-Python wheels) and the full pip dependency tree into a virtualenv. The runtime stage starts fresh from a clean python:3.11-slim base and copies only the finished virtualenv and application code across — none of the compilers or apt package caches make it into the final image. The result is typically 150–250MB smaller than a single-stage equivalent, and a smaller image means a smaller attack surface and faster pulls when scaling out.

Building and Running

# Build the image, tagging it with both a semantic version and 'latest'
docker build -t churn-api:1.0.0 -t churn-api:latest .

# Run it, mapping the container's port 8000 to the host, injecting secrets
# via environment variables rather than baking them into the image
docker run -d \
    --name churn-api \
    -p 8000:8000 \
    -e API_KEY=supersecretkey123 \
    -e MODEL_SOURCE=local \
    churn-api:latest

# Confirm it's healthy
curl http://localhost:8000/health

7 docker-compose: API + MLflow for Local Testing

In Lesson 65 you ran an MLflow tracking server to log experiments and register models. To test the "pull from registry" code path end-to-end before deploying anywhere, you want the API container and the MLflow server running together, networked so the API can reach MLflow by service name. docker-compose is the standard tool for exactly this: declare every service your stack needs in one YAML file and bring the whole thing up with one command.

# docker-compose.yml
version: "3.9"

services:
  mlflow:
    image: ghcr.io/mlflow/mlflow:v2.14.1
    command: >
      mlflow server
      --backend-store-uri sqlite:///mlflow/mlflow.db
      --default-artifact-root /mlflow/artifacts
      --host 0.0.0.0
      --port 5000
    ports:
      - "5000:5000"
    volumes:
      - mlflow-data:/mlflow
    healthcheck:
      test: ["CMD", "curl", "-f", "http://localhost:5000/health"]
      interval: 10s
      timeout: 5s
      retries: 5

  churn-api:
    build: .
    image: churn-api:latest
    ports:
      - "8000:8000"
    environment:
      MODEL_SOURCE: "mlflow"
      MLFLOW_TRACKING_URI: "http://mlflow:5000"
      MLFLOW_MODEL_NAME: "churn-classifier"
      MLFLOW_MODEL_STAGE: "Production"
      API_KEY: "${API_KEY:?Set API_KEY in your .env file}"
    depends_on:
      mlflow:
        condition: service_healthy
    restart: on-failure

volumes:
  mlflow-data:

Note MLFLOW_TRACKING_URI: "http://mlflow:5000" — inside the compose network, services address each other by their service name (mlflow), not localhost. This is the same DNS-by-service-name pattern Kubernetes uses, so testing it locally with compose is good practice for the real deployment topology in Lesson 67.

# Bring up the whole stack
docker compose up --build -d

# Watch logs from both services
docker compose logs -f churn-api

# Tear down (and remove the named volume to fully reset MLflow's state)
docker compose down -v

8 Load Testing Before You Ship

A model that returns correct predictions but takes 4 seconds per request, or that falls over at 50 concurrent users, is not production-ready. Before handing the service to an SRE team or pointing real traffic at it, measure its throughput and latency under load. Locust is a popular Python-based load-testing tool: you describe user behavior as Python code, and it simulates many concurrent "users" hammering your API while reporting response time percentiles in real time.

# locustfile.py
from locust import HttpUser, task, between
import random

class ChurnAPIUser(HttpUser):
    wait_time = between(0.1, 0.5)   # think time between requests, seconds

    def on_start(self):
        self.headers = {"X-API-Key": "supersecretkey123"}

    @task(3)
    def predict_single(self):
        payload = {
            "tenure_months": random.randint(0, 72),
            "monthly_charges": round(random.uniform(20, 150), 2),
            "total_charges": round(random.uniform(20, 8000), 2),
            "contract_type": random.choice(["month-to-month", "one-year", "two-year"]),
            "num_support_tickets": random.randint(0, 10),
            "has_autopay": random.choice([True, False]),
        }
        self.client.post("/v1/predict", json=payload, headers=self.headers)

    @task(1)
    def health_check(self):
        self.client.get("/health")

# Run a headless load test: 200 simulated users, ramping up at 10/sec, for 2 minutes
locust -f locustfile.py --host http://localhost:8000 \
    --headless -u 200 -r 10 --run-time 2m --csv=results

The numbers tell you whether the service meets its target SLA before a real outage tells you. If p99 latency balloons under load, common culprits are too few uvicorn workers for the available CPU cores, a model that's slower than it needs to be (consider quantization or a lighter model), or synchronous I/O calls blocking the event loop where they shouldn't be. If error rate climbs instead, check for resource exhaustion (memory, file descriptors) or unhandled exceptions surfacing as 500s under concurrent load that never appeared in single-request testing.

The single Locust run above reports percentiles at one concurrency level (200 simulated users). To decide where a service actually starts to struggle, run the same test repeatedly at increasing concurrency and plot how each percentile moves — flat lines mean you have headroom, and a sharp upward bend marks the point where the deployment (here, 2 uvicorn workers per replica) runs out of CPU to keep up.

9 Making the Model Itself Faster: Quantization, Pruning & Distillation

Section 8's load test tuned infrastructure — workers, concurrency, container resources — to serve a model faster. But often the biggest latency and cost win comes from making the model smaller and cheaper to run in the first place, before it ever reaches Docker. Three techniques dominate in practice, and they compose with each other.

Quantization: Fewer Bits Per Number

A typical PyTorch model stores every weight as a 32-bit float. Quantization converts those weights (and often activations) to a lower-precision format — 8-bit integers, or even 4-bit for large language models — trading a small amount of numerical precision for a large reduction in memory footprint and inference latency.

import torch
from torch.ao.quantization import quantize_dynamic
import torch.nn as nn

model = nn.Sequential(
    nn.Linear(512, 256), nn.ReLU(),
    nn.Linear(256, 128), nn.ReLU(),
    nn.Linear(128, 10),
)
model.eval()

# Dynamic quantization: weights converted to int8 ahead of time, activations
# quantized on-the-fly during inference -- the simplest form, no retraining needed.
quantized_model = quantize_dynamic(model, {nn.Linear}, dtype=torch.qint8)

def model_size_mb(m):
    torch.save(m.state_dict(), '/tmp/_tmp_model.pt')
    import os
    size = os.path.getsize('/tmp/_tmp_model.pt') / 1e6
    os.remove('/tmp/_tmp_model.pt')
    return size

print(f"FP32 model size:  {model_size_mb(model):.2f} MB")
print(f"INT8 model size:  {model_size_mb(quantized_model):.2f} MB")
# Typically ~4x smaller (32 bits -> 8 bits), with inference speedups on CPU
# that are often even larger due to better memory bandwidth utilization.

Pruning: Removing Unimportant Weights

Pruning exploits the fact that large trained networks are typically over-parameterized — many weights contribute almost nothing to the output and can be zeroed out (or removed entirely) with minimal accuracy loss.

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

layer = nn.Linear(512, 256)

# Unstructured magnitude pruning: zero out the 30% smallest-magnitude weights
prune.l1_unstructured(layer, name='weight', amount=0.3)

remaining = (layer.weight != 0).float().mean().item()
print(f"Fraction of weights remaining after pruning: {remaining:.1%}")

# Unstructured pruning creates SPARSE weight matrices -- real speedup requires
# either specialized sparse-matrix hardware/kernels, or STRUCTURED pruning,
# which removes entire neurons/channels so the resulting matrix is simply
# smaller and dense, and any framework benefits immediately:
# prune.ln_structured(layer, name='weight', amount=0.3, n=2, dim=0)

Knowledge Distillation: Training a Small Model to Mimic a Large One

Distillation (Hinton et al., 2015) takes a different approach entirely: instead of shrinking a trained model, train a brand-new, much smaller "student" model to reproduce a large, accurate "teacher" model's output distribution — not just its hard predicted labels, but its full soft probabilities, which encode extra information about how confident and how similar different classes are (this is the same teacher-student framing behind DistilBERT, a 40%-smaller BERT that retains ~97% of its performance).

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

def distillation_loss(student_logits, teacher_logits, true_labels, temperature=3.0, alpha=0.5):
    # Soften both distributions with temperature -- higher T reveals more
    # of the teacher's relative confidence across ALL classes, not just the winner
    soft_teacher = F.softmax(teacher_logits / temperature, dim=1)
    soft_student = F.log_softmax(student_logits / temperature, dim=1)
    distill_loss = F.kl_div(soft_student, soft_teacher, reduction='batchmean') * (temperature ** 2)

    # Still anchor to the real labels too, so the student doesn't ONLY imitate the teacher
    hard_loss = F.cross_entropy(student_logits, true_labels)

    return alpha * distill_loss + (1 - alpha) * hard_loss

# teacher_model.eval()  # frozen, no gradients
# for x, y in train_loader:
#     with torch.no_grad():
#         teacher_logits = teacher_model(x)
#     student_logits = student_model(x)
#     loss = distillation_loss(student_logits, teacher_logits, y)
#     loss.backward(); optimizer.step()

Exporting for a Portable, Optimized Runtime: ONNX

Once a model is trained (and optionally quantized/pruned/distilled), exporting it to ONNX (Open Neural Network Exchange) decouples it from PyTorch's Python runtime entirely, letting it run through onnxruntime — a C++ inference engine with graph-level optimizations (operator fusion, constant folding) that frequently outperforms native PyTorch inference, with no Python interpreter overhead in the serving path at all.

import torch
import onnxruntime as ort
import numpy as np

model.eval()
dummy_input = torch.randn(1, 512)

torch.onnx.export(
    model, dummy_input, "model.onnx",
    input_names=['input'], output_names=['output'],
    dynamic_axes={'input': {0: 'batch_size'}, 'output': {0: 'batch_size'}},  # variable batch size
    opset_version=17,
)

# Serve with onnxruntime instead of PyTorch -- this is what would sit
# inside the FastAPI /predict endpoint from Section 3
session = ort.InferenceSession("model.onnx", providers=['CPUExecutionProvider'])
input_data = np.random.randn(1, 512).astype(np.float32)
outputs = session.run(None, {'input': input_data})
print(f"ONNX Runtime output shape: {outputs[0].shape}")