FastAPI 部署深度学习模型的方法

安装依赖

确保安装 FastAPI 和 Uvicorn：

pip install fastapi uvicorn

如果模型依赖其他库（如 PyTorch、TensorFlow），需一并安装：

pip install torch tensorflow

创建 FastAPI 应用

新建一个 Python 文件（如 main.py），初始化 FastAPI 应用：

from fastapi import FastAPI
app = FastAPI()

加载模型

在应用启动时加载预训练模型。以 PyTorch 为例：

import torch
from transformers import pipeline

model = pipeline("text-classification", model="distilbert-base-uncased-finetuned-sst-2-english")

定义预测接口

创建一个 POST 接口接收输入数据并返回预测结果：

from pydantic import BaseModel

class InputText(BaseModel):
    text: str

@app.post("/predict")
def predict(input: InputText):
    prediction = model(input.text)
    return {"prediction": prediction}

运行服务

使用 Uvicorn 启动服务：

uvicorn main:app --reload

访问 http://127.0.0.1:8000/docs 可查看交互式 API 文档。

异步支持（可选）

对于高并发场景，可以使用异步方式加载模型：

from fastapi import FastAPI
import asyncio

app = FastAPI()

async def load_model():
    return pipeline("text-classification", model="distilbert-base-uncased-finetuned-sst-2-english")

model = asyncio.run(load_model())

生产环境部署

使用 Gunicorn 管理多进程：

pip install gunicorn
gunicorn -w 4 -k uvicorn.workers.UvicornWorker main:app

Docker 容器化

创建 Dockerfile：

FROM python:3.9
WORKDIR /app
COPY requirements.txt .
RUN pip install -r requirements.txt
COPY . .
CMD ["gunicorn", "-w", "4", "-k", "uvicorn.workers.UvicornWorker", "main:app"]

构建并运行容器：

docker build -t model-api .
docker run -p 8000:8000 model-api