训练好的 AI 模型怎么上线？FastAPI、ONNX、Docker 部署教程

完整链路：训练 → 导出模型 → 封装 HTTP 服务 → 容器化 → 部署到服务器。

第 1 步：保存训练好的模型

# PyTorch 推荐只存参数
torch.save(model.state_dict(), "model.pt")
# 如需跨框架部署，转 ONNX（可选但推荐）
torch.onnx.export(model, dummy_input, "model.onnx", input_names=["input"], output_names=["output"])

第 2 步：用 FastAPI 封装成 API（最主流）

from fastapi import FastAPI
from pydantic import BaseModel
import torch

app = FastAPI()
model = torch.load("model.pt", map_location="cpu")  # 服务启动时加载一次
model.eval()

class Item(BaseModel):
    features: list[float]

@app.post("/predict")
def predict(item: Item):
    x = torch.tensor([item.features])
    with torch.no_grad():
        pred = model(x).tolist()[0]
    return {"prediction": pred}

启动：uvicorn app:app --host 0.0.0.0 --port 8000。自带 /docs 交互式文档，方便测试。注意：模型只加载一次，不要在请求里反复加载。

第 3 步：Docker 容器化（解决"在我电脑上能跑"）

FROM python:3.11-slim
WORKDIR /app
COPY requirements.txt .
RUN pip install -i https://pypi.tuna.tsinghua.edu.cn/simple -r requirements.txt
COPY . .
EXPOSE 8000
CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "8000"]

docker build -t my-ai-service .
docker run -p 8000:8000 my-ai-service

第 4 步：上线与运维

• 有 GPU 用 GPU 镜像（nvidia/cuda 基础镜像），部署前确认 torch.cuda.is_available()。
• 生产考虑：nginx 反代 + 多 worker（uvicorn --workers 4）、模型服务化工具（TorchServe、NVIDIA Triton）、健康检查 /health、日志与监控（Prometheus）。
• 升级更新：灰度发布（先切 10% 流量验证再全量）。

新手最小可行路径： FastAPI 起服务（本机验证）→ Docker 打包（到处能跑）→ 云服务器跑起来 → 再加监控与并发优化。