这篇的目标不是训练模型,而是把一个现成模型包装成生产可用的 API。这是 AI 应用开发里最高频的需求,难点不在模型,在工程:异步、并发限制、错误处理、结果缓存。
一、选型:用现成模型,别自己训
| 方案 | 适合 | 判断 |
|---|---|---|
| ONNX Runtime + 现成模型 | CPU 推理、跨平台 | 本项目选它:无需 GPU,镜像小 |
| PyTorch 直接加载 | GPU 推理、可微调 | 镜像 2-4GB,CPU 上慢 |
| 调用云端视觉 API | 不想部署模型 | 按次计费,有网络依赖 |
用 onnxruntime 加载一个预训练的图像分类模型(如 ResNet50 / MobileNetV3),CPU 上单张推理 10-50ms,足够用。
二、项目结构
api/
├── app/
│ ├── main.py # FastAPI 入口
│ ├── model.py # 模型加载与推理
│ ├── schemas.py # 请求/响应模型
│ └── deps.py # 依赖注入
├── models/
│ └── mobilenetv3.onnx # 预训练模型
├── requirements.txt
└── Dockerfile三、模型加载:启动时一次,别每请求加载
# app/model.py
import onnxruntime as ort
import numpy as np
from PIL import Image
from functools import lru_cache
from pathlib import Path
MODEL_PATH = Path(__file__).parent.parent / "models" / "mobilenetv3.onnx"
class Classifier:
def __init__(self, path: Path):
1. CPU 上开满线程;intra_op_num_threads 控制单次推理并行度
so = ort.SessionOptions()
so.intra_op_num_threads = 4
so.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
self.sess = ort.InferenceSession(str(path), so, providers=["CPUExecutionProvider"])
self.input_name = self.sess.get_inputs()[0].name
def preprocess(self, img: Image.Image) -> np.ndarray:
1. ImageNet 标准预处理
img = img.convert("RGB").resize((224, 224))
x = np.asarray(img, dtype=np.float32) / 255.0
mean = np.array([0.485, 0.456, 0.406], dtype=np.float32)
std = np.array([0.229, 0.224, 0.225], dtype=np.float32)
x = (x - mean) / std
return x.transpose(2, 0, 1)[None, ...] # NCHW
def predict(self, img: Image.Image) -> tuple[int, float]:
x = self.preprocess(img)
out = self.sess.run(None, {self.input_name: x})[0][0]
1. softmax
e = np.exp(out - out.max())
p = e / e.sum()
idx = int(p.argmax())
return idx, float(p[idx])
@lru_cache(maxsize=1)
def get_classifier() -> Classifier:
"""全局单例,首次调用时加载。"""
return Classifier(MODEL_PATH)- 模型加载放
lru_cache单例:加载一次约几百毫秒,每请求加载会直接压垮服务 intra_op_num_threads别设太大:超过 CPU 核数会反而变慢(线程争抢)- 预处理必须与训练时一致:mean/std 的顺序(RGB 还是 BGR)错了,准确率会莫名下降
四、API 层:异步 + 并发控制
# app/main.py
from fastapi import FastAPI, File, UploadFile, HTTPException
from fastapi.concurrency import run_in_threadpool
from PIL import Image
import io, asyncio
from app.model import get_classifier, LABELS
app = FastAPI(title="Image Classifier API")
1. 限制并发推理数:超过 CPU 核数没意义,反而排队变慢
_sem = asyncio.Semaphore(4)
@app.post("/predict")
async def predict(file: UploadFile = File(...)):
1. 1. 校验类型与大小
if file.content_type not in {"image/jpeg", "image/png", "image/webp"}:
raise HTTPException(415, "只支持 jpeg/png/webp")
raw = await file.read()
if len(raw) > 5 * 1024 * 1024:
raise HTTPException(413, "图片不能超过 5MB")
1. 2. 解码(在线程池里做,避免阻塞事件循环)
try:
img = await run_in_threadpool(Image.open, io.BytesIO(raw))
img.load()
except Exception:
raise HTTPException(400, "图片解析失败")
1. 3. 推理(同样是 CPU 密集,必须放线程池)
async with _sem:
cls = get_classifier() # 单例
idx, conf = await run_in_threadpool(cls.predict, img)
return {
"label": LABELS[idx],
"confidence": round(conf, 4),
}两个关键点:① 所有 CPU 密集操作必须用
run_in_threadpool,否则会阻塞 asyncio 事件循环,整个服务卡住;② 用Semaphore限制并发推理数,超过 CPU 核数只会让排队更慢。
五、常见错误处理对照
| 情况 | 状态码 | 处理 |
|---|---|---|
| 非图片类型 | 415 | 校验 content_type |
| 文件过大 | 413 | 读之前/后检查 size |
| 图片损坏 | 400 | Image.open + load 捕获异常 |
| 模型未加载 | 503 | 启动时校验模型文件存在 |
| 并发过高 | 429 | Semaphore + 超时返回 429 |
六、加了缓存后 QPS 提升明显
很多用户会传同一张图(重试、批量)。用图片内容哈希做缓存,能省掉大量重复推理:
import hashlib, json
import redis.asyncio as redis
rdb = redis.from_url("redis://redis:6379/0")
@app.post("/predict")
async def predict(file: UploadFile = File(...)):
raw = await file.read()
key = "pred:" + hashlib.sha256(raw).hexdigest()[:32]
1. 命中缓存直接返回
if cached := await rdb.get(key):
return json.loads(cached)
img = await run_in_threadpool(Image.open, io.BytesIO(raw))
img.load()
async with _sem:
cls = get_classifier()
idx, conf = await run_in_threadpool(cls.predict, img)
result = {"label": LABELS[idx], "confidence": round(conf, 4)}
await rdb.setex(key, 3600, json.dumps(result))
return result七、部署:镜像要小
FROM python:3.12-slim
WORKDIR /app
1. 只装 CPU 版 onnxruntime(不要装 onnxruntime-gpu,镜像会大 10 倍)
RUN pip install --no-cache-dir \
fastapi uvicorn[standard] onnxruntime pillow python-multipart redis
COPY app/ ./app/
COPY models/ ./models/
1. 非 root 运行
RUN useradd -m appuser && chown -R appuser /app
USER appuser
EXPOSE 8000
CMD ["uvicorn", "app.main:app", "--host", "0.0.0.0", "--port", "8000", "--workers", "1"]--workers 1是有意的:多 worker 会在每个进程各加载一份模型,内存翻倍- 要提并发就调
intra_op_num_threads与 Semaphore,而不是加 worker - 不要装 onnxruntime-gpu:除非真有 GPU,否则镜像体积大 10 倍还跑不起来
八、压测与真实数字
# 准备测试图
curl -o test.jpg https://example.com/sample.jpg
1. 压测(注意用 multipart 上传)
hey -n 200 -c 8 -m POST \
-F "file=@test.jpg" \
http://localhost:8000/predict| 配置 | QPS | P99 | 说明 |
|---|---|---|---|
| 2C 无缓存 | 约 60 | 120ms | 推理占满 CPU |
| 2C + 缓存(重复图) | 约 2000 | 8ms | 命中缓存直接返回 |
| 4C 无缓存 | 约 110 | 90ms | 核数接近线性 |
| 4C + 缓存 | 约 3800 | 6ms | 受限于网络与 JSON 序列化 |
结论:缓存的收益远大于加核数。真实场景里重复图片的比例很高(用户重试、批量导图),所以缓存是这类 API 的第一优化项。
九、这个项目教会了什么
- 模型加载必须是单例:每请求加载会直接打死服务
- CPU 密集必须丢线程池:在 async 函数里直接跑推理会阻塞整个事件循环
- 并发控制比加核数重要:Semaphore 限制在核数附近,多了只是排队
- 缓存是这类 API 的最大优化点:内容哈希去重,收益远超硬件升级
- 镜像选 CPU 版依赖:onnxruntime-gpu 会让镜像大 10 倍且无法在无 GPU 机器运行