📚 课程目录第3章使用FastAPI实现高性能API

3. 使用FastAPI实现高性能API

任务目标

构建一个基于FastAPI的异步图像分类推理服务，使用Pydantic进行请求体验证，并通过异步处理支持并发请求。客户端提交图像URL，服务端下载图像并返回分类结果。

环境依赖

pip install fastapi uvicorn pillow torch torchvision httpx python-multipart

完整实现代码

import asyncio
from io import BytesIO
from typing import Optional

import httpx
import torch
import torchvision.transforms as transforms
from fastapi import FastAPI, HTTPException
from PIL import Image
from pydantic import BaseModel, HttpUrl, validator

# ---------- 请求/响应模型 ----------
class InferenceRequest(BaseModel):
    image_url: HttpUrl
    top_k: Optional[int] = 3

    @validator("top_k")
    def validate_top_k(cls, v):
        if v is not None and (v < 1 or v > 10):
            raise ValueError("top_k must be between 1 and 10")
        return v

class Prediction(BaseModel):
    label: str
    score: float

class InferenceResponse(BaseModel):
    predictions: list[Prediction]

# ---------- 模型加载 ----------
model = None
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

# ImageNet类别标签（简化示例，仅前5类用于演示）
IMAGENET_LABELS = ["tench", "goldfish", "great white shark", "tiger shark", "hammerhead shark"]

def load_model():
    global model
    # 使用预训练ResNet18，仅用于演示推理流程
    model = torch.hub.load("pytorch/vision:v0.10.0", "resnet18", pretrained=True)
    model.eval()

# ---------- 图像下载与预处理 ----------
async def download_image(url: str) -> Image.Image:
    async with httpx.AsyncClient(timeout=10) as client:
        resp = await client.get(url)
        if resp.status_code != 200:
            raise HTTPException(status_code=400, detail="Failed to download image")
        return Image.open(BytesIO(resp.content)).convert("RGB")

def preprocess(image: Image.Image) -> torch.Tensor:
    return transform(image).unsqueeze(0)

# ---------- 推理 ----------
async def predict(image: Image.Image, top_k: int) -> list[Prediction]:
    # 在线程池中执行CPU/GPU推理，避免阻塞事件循环
    loop = asyncio.get_event_loop()
    tensor = preprocess(image)
    with torch.no_grad():
        logits = await loop.run_in_executor(None, lambda: model(tensor))
    probs = torch.softmax(logits, dim=1).squeeze(0)
    top_probs, top_indices = torch.topk(probs, k=top_k)
    return [
        Prediction(label=IMAGENET_LABELS[idx], score=float(prob))
        for prob, idx in zip(top_probs, top_indices)
    ]

# ---------- FastAPI应用 ----------
app = FastAPI(title="Image Classification API")

@app.on_event("startup")
async def startup():
    load_model()

@app.post("/predict", response_model=InferenceResponse)
async def inference(req: InferenceRequest):
    """接收图像URL，返回top_k分类结果"""
    image = await download_image(str(req.image_url))
    predictions = await predict(image, req.top_k or 3)
    return InferenceResponse(predictions=predictions)

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

运行方式

uvicorn main:app --host 0.0.0.0 --port 8000 --workers 4

验证方法

启动服务：执行上述命令，看到“Application startup complete”日志。
健康检查：访问 http://localhost:8000/health，应返回 {"status":"ok"}。
推理请求：使用curl发送POST请求：

curl -X POST http://localhost:8000/predict \
  -H "Content-Type: application/json" \
  -d '{"image_url": "https://upload.wikimedia.org/wikipedia/commons/3/3a/Cat03.jpg", "top_k": 3}'

通过条件：返回JSON包含predictions数组，每个元素含label和score字段，且score值在0~1之间。

并发验证

# 使用ab工具发送20个并发请求
ab -n 20 -c 10 -p request.json -T application/json http://localhost:8000/predict

观察所有请求均返回200状态码，无超时或连接错误，说明异步处理生效。

关键说明

模型加载在startup事件中完成，避免首次请求延迟。
推理通过run_in_executor放入线程池执行，防止阻塞事件循环，保证并发能力。
Pydantic自动验证请求体：image_url必须是合法URL，top_k范围限制在1~10，非法输入返回422错误。
示例仅加载前5类ImageNet标签，完整标签文件需替换IMAGENET_LABELS为实际1000类列表。