Python AI 实战（18/20）：训练好的 AI 模型怎么上线？FastAPI、ONNX、Docker 部署教程

Python AI 实战问答速查手册 · 第 18/20 问

所属板块：部署上线

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

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

# PyTorch 推荐只存参数
torch.save(model.state_dict(), "model.pt")
# model 已完成训练；先切到评估模式
model.eval()
# 可选：dummy_input 需与真实输入的形状和类型匹配，导出后核对输出一致性
torch.onnx.export(model, dummy_input, "model.onnx", input_names=["input"], output_names=["output"])

第 2 步：用 FastAPI 封装成 API 下例假设训练时的模型类保存在 model_def.py 中，名称为 MyModel，接收两个浮点特征。请替换为实际模型和训练时完全一致的结构、初始化参数与预处理。

from fastapi import FastAPI
from pydantic import BaseModel, Field, FiniteFloat
import torch
from model_def import MyModel  # 必须与训练时的模型定义一致

app = FastAPI()
model = MyModel()
state = torch.load("model.pt", map_location="cpu", weights_only=True)
model.load_state_dict(state)
model.eval()

class Item(BaseModel):
  features: list[FiniteFloat] = Field(min_length=2, max_length=2)

@app.get("/health")
def health():
  return {"status": "ok"}

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

启动：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 部署需兼容的运行环境，并显式把模型和输入移到同一设备。上面的服务示例使用 CPU。

• 生产考虑：nginx 反代、鉴权和限流、健康检查、日志与监控。先单 worker 压测，每个 worker 通常会各自加载模型，增加 worker 前评估内存和显存；需要专用服务引擎时再评估 NVIDIA Triton 等工具。

• 升级更新：灰度发布（先切 10% 流量验证再全量）。

新手最小可行路径： FastAPI 起服务（本机验证）→ Docker 打包（在兼容的硬件和运行环境中验证）→ 云服务器跑起来 → 再加监控与并发优化。