TorchServe 与 ONNX：模型部署框架

把训练好的模型变成可以调用的 API 服务。

1.1 为什么需要专门的部署框架

TorchServe vs ONNX Runtime对比图
TorchServe 与 ONNX Runtime 在部署流程、支持格式、性能及适用场景的全面对比

直接用 Flask/FastAPI 部署模型的问题是什么？

先说清楚，对于简单场景，FastAPI 直接加载模型完全够用。但当你的需求超出简单场景时，就会遇到这些问题：

没有批处理优化：一个请求一个请求地处理，GPU 利用率很低。用批处理同时处理多个请求，吞吐量可以提升 2-5 倍。
没有模型版本管理：同时跑 A/B 测试的两个模型版本，怎么管？旧版怎么快速回滚？
没有健康检查：服务出问题了，负载均衡器怎么知道不要再往这台机器转流量？
扩展性差：单机撑不住，怎么部署多实例？请求怎么路由？

TorchServe 是 PyTorch 官方提供的解决方案，解决了上面这些问题。

1.2 场景选型：什么时候用什么

场景	推荐方案	理由
快速原型验证	FastAPI + 直接加载模型	最简单，几行代码搞定
中等规模生产	FastAPI + Docker + Nginx	熟悉的栈，容易维护
LLM 推理服务	vLLM（见第 04 篇）	专为 LLM 优化，10-30x 吞吐提升
传统 ML 模型批量服务	TorchServe	批处理、版本管理、多模型
无 GPU 的 CPU 推理	ONNX Runtime	CPU 上比 PyTorch 快 1.5-3x

选型原则：先考虑 FastAPI，遇到性能瓶颈或需要高级功能时再换专业框架。不要过度工程化。

1.3 TorchServe：PyTorch官方部署框架

bash

pip install torchserve torch-model-archiver

1.3.1 步骤1：创建Handler

python

# handler.py
import torch
from ts.torch_handler.base_handler import BaseHandler
from transformers import AutoTokenizer, AutoModelForSequenceClassification

class SentimentHandler(BaseHandler):
    def initialize(self, context):
        """加载模型（只执行一次）"""
        self.model_dir = context.system_properties.get("model_dir")
        self.tokenizer = AutoTokenizer.from_pretrained(self.model_dir)
        self.model = AutoModelForSequenceClassification.from_pretrained(self.model_dir)
        self.model.eval()
        self.initialized = True
    
    def preprocess(self, data):
        """预处理输入"""
        texts = [item.get("data") or item.get("body") for item in data]
        inputs = self.tokenizer(
            texts, return_tensors="pt",
            padding=True, truncation=True, max_length=512
        )
        return inputs
    
    def inference(self, inputs):
        """模型推理"""
        with torch.no_grad():
            outputs = self.model(**inputs)
        return outputs.logits
    
    def postprocess(self, logits):
        """后处理输出"""
        predictions = torch.softmax(logits, dim=-1)
        results = []
        for pred in predictions:
            label_id = pred.argmax().item()
            confidence = pred.max().item()
            label = self.model.config.id2label[label_id]
            results.append({
                "label": label,
                "confidence": round(confidence, 4)
            })
        return results

1.3.2 步骤2：打包模型

bash

torch-model-archiver \
    --model-name sentiment-classifier \
    --version 1.0 \
    --model-file model/ \
    --handler handler.py \
    --export-path model-store/

1.3.3 步骤3：启动服务

bash

torchserve --start \
    --model-store model-store \
    --models sentiment=sentiment-classifier.mar \
    --ncs

# 调用API
curl -X POST http://localhost:8080/predictions/sentiment \
    -H "Content-Type: application/json" \
    -d '{"data": "这个产品真的太棒了！"}'

1.4 FastAPI + 模型服务（轻量方案）

对于中小规模应用，FastAPI + 简单的模型加载也是不错的选择：

python

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from transformers import pipeline
import torch
from typing import List
import asyncio
from concurrent.futures import ThreadPoolExecutor

app = FastAPI(title="AI模型服务")

# 全局模型（只加载一次）
model = None
executor = ThreadPoolExecutor(max_workers=4)

@app.on_event("startup")
async def load_model():
    global model
    model = pipeline(
        "text-classification",
        model="bert-base-chinese",
        device=0 if torch.cuda.is_available() else -1
    )
    print("模型加载完成")

class PredictRequest(BaseModel):
    texts: List[str]

class PredictResponse(BaseModel):
    results: List[dict]
    model_version: str = "1.0"

@app.post("/predict", response_model=PredictResponse)
async def predict(request: PredictRequest):
    if model is None:
        raise HTTPException(status_code=503, detail="模型未加载")
    
    # 在线程池中运行CPU密集型任务
    loop = asyncio.get_event_loop()
    results = await loop.run_in_executor(
        executor, model, request.texts
    )
    
    return PredictResponse(results=results)

@app.get("/health")
async def health_check():
    return {"status": "healthy", "model_loaded": model is not None}

# 运行: uvicorn main:app --host 0.0.0.0 --port 8000 --workers 4

1.5 容器化部署

容器化（Containerization）：将应用及其所有依赖打包成一个独立的"容器"，在任何机器上运行结果都一致，不再出现"我这里能跑"的问题。Docker 是最常用的容器化工具。

dockerfile

# Dockerfile
FROM pytorch/pytorch:2.1.0-cuda12.1-cudnn8-runtime

WORKDIR /app

COPY requirements.txt .
RUN pip install -r requirements.txt

COPY . .

# 预下载模型（构建时缓存）
RUN python -c "from transformers import AutoTokenizer, AutoModel; AutoTokenizer.from_pretrained('bert-base-chinese')"

EXPOSE 8000

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

yaml

# docker-compose.yml
version: '3.8'
services:
  ai-service:
    build: .
    ports:
      - "8000:8000"
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              count: 1
              capabilities: [gpu]
    environment:
      - MODEL_NAME=bert-base-chinese
    restart: unless-stopped
  
  nginx:
    image: nginx:latest
    ports:
      - "80:80"
    volumes:
      - ./nginx.conf:/etc/nginx/nginx.conf
    depends_on:
      - ai-service

1.6 小结

场景	推荐方案
快速原型	FastAPI + 直接加载模型
中等规模	FastAPI + Docker + Nginx
大规模生产	TorchServe / vLLM
LLM推理	vLLM（见下一篇）

选择部署方案时，先考虑需求规模，不要过度工程化。