使用 FastAPI 与 ONNX 部署机器学习模型的最佳实践

import tensorflow as tf
import tf2onnx

# 加载训练好的模型
model = tf.keras.models.load_model('model.h5')

# 转换为 ONNX 格式
spec = (tf.TensorSpec(model.inputs[0].shape, dtype=tf.float32, name="input"),)
output_path = "model.onnx"
model_proto, _ = tf2onnx.convert.from_keras(model, 
                                            input_signature=spec, 
                                            opset=11, 
                                            output_path=output_path)

print("✅ 模型已成功转换为 ONNX 格式！")

验证转换后的模型

转换完成后，别忘了验证一下模型是否正常工作！

import onnx
import onnxruntime as ort

# 加载 ONNX 模型
onnx_model = onnx.load('model.onnx')
onnx.checker.check_model(onnx_model)
print("✅ 模型格式验证通过！")

# 使用 ONNX Runtime 进行推理
ort_session = ort.InferenceSession('model.onnx')

# 准备输入数据
import numpy as np
input_data = np.random.randn(1, 3, 224, 224).astype(np.float32)

# 运行推理
outputs = ort_session.run(None, {'input': input_data})

print("输出结果：", outputs)

2. 构建 FastAPI 应用

现在，我们来创建一个基于 FastAPI 的应用，将模型部署为一个 API 服务。

安装依赖

首先，安装必要的依赖包：

pip install fastapi uvicorn[standard] onnxruntime

定义 FastAPI 应用

编写应用主文件 main.py：

from fastapi import FastAPI
import onnxruntime as ort
import numpy as np
from pydantic import BaseModel

app = FastAPI(title="机器学习模型部署 API")

# 加载 ONNX 模型
ort_session = ort.InferenceSession('model.onnx')

# 定义输入数据模型
class InputData(BaseModel):
    data: list

@app.post("/predict")
async def predict(input_data: InputData):
    # 将输入数据转换为 numpy 数组
    input_array = np.array(input_data.data).astype(np.float32)
    
    # 进行推理
    outputs = ort_session.run(None, {'input': input_array})
    
    # 返回结果
    return {"prediction": outputs[0].tolist()}

运行应用

使用 uvicorn 启动应用：

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

测试接口

可以使用 curl 或其他工具测试一下接口是否正常工作：

curl -X POST "http://localhost:8000/predict" -H "Content-Type: application/json" -d '{
    "data": [0.5, 0.3, 0.2]
}'

示例输出：

{
    "prediction": [[0.1, 0.9]]
}

至此我们的 API 已经可以正常工作了！

3. 性能优化

性能对于一个服务来说至关重要，这里介绍一些优化技巧。

模型优化

• 使用模型优化工具：ONNX 提供了模型优化工具，可简化和加速模型。

python -m onnxruntime.tools.optimizer_cli --input model.onnx --output model_optimized.onnx --optimization_level all

• 量化模型：通过模型量化，将浮点数精度降低，减小模型大小，加速推理。

python -m onnxruntime.quantization.quantize --input model.onnx --output model_quant.onnx --per_channel

推理加速

• 使用 GPU 加速：如果有 GPU 资源，可以使用 GPU 提供商提升推理速度。

ort_session = ort.InferenceSession('model.onnx', providers=['CUDAExecutionProvider'])

• 多线程或多进程：根据服务器性能，调整并发数，充分利用硬件资源。

4. 安全性考虑

安全是服务的底线，我们需要考虑以下几点。

输入验证

• 数据格式验证：使用 Pydantic 模型，确保输入数据的格式和类型正确。
• 异常处理：捕获可能的异常，如数据维度错误，返回友好的错误信息。

@app.post("/predict")
async def predict(input_data: InputData):
    try:
        # 输入验证
        input_array = np.array(input_data.data).astype(np.float32)
        # 检查输入维度（根据模型需求调整）
        if input_array.shape != (1, 3, 224, 224):
            return {"error": "输入数据维度不正确"}
        # 进行推理
        outputs = ort_session.run(None, {'input': input_array})
        return {"prediction": outputs[0].tolist()}
    except Exception as e:
        return {"error": str(e)}

安全防护

• 限制请求频率：通过中间件或网关，防止恶意请求和 DDoS 攻击。
• SSL/HTTPS：在生产环境中，确保通信的安全性。

5. 容器化部署

为了便于在不同环境中部署和扩展，推荐使用 Docker 容器化 FastAPI 应用。

编写 Dockerfile

创建 Dockerfile 文件：

FROM python:3.9-slim

WORKDIR /app

COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt

COPY . .

EXPOSE 8000

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

构建与运行

docker build -t ml-api-app .
docker run -p 8000:8000 ml-api-app

容器化部署确保了环境的一致性，避免了'在我机器上能跑'的问题，同时便于与 Kubernetes 等编排工具集成。

案例示例

下面以一个手写数字识别模型为例，展示完整的部署过程。

1. 模型训练与转换

# 使用 MNIST 数据集训练模型
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms

# 定义简单的神经网络
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.fc = nn.Linear(28*28, 10)
        
    def forward(self, x):
        x = x.view(-1, 28*28)
        x = self.fc(x)
        return x

# 创建模型实例
model = Net()
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.01)

# 加载数据集
train_loader = torch.utils.data.DataLoader(
    datasets.MNIST('.', train=True, download=True, transform=transforms.ToTensor()),
    batch_size=64, shuffle=True)

# 训练模型（这里只训练一个 epoch 作示例）
for epoch in range(1):
    for data, target in train_loader:
        optimizer.zero_grad()
        output = model(data)
        loss = criterion(output, target)
        loss.backward()
        optimizer.step()

# 保存模型
torch.save(model.state_dict(), 'mnist.pth')

print("✅ 模型训练完成并已保存！")

# 加载模型并转换为 ONNX
model = Net()
model.load_state_dict(torch.load('mnist.pth'))
model.eval()

dummy_input = torch.randn(1, 1, 28, 28)
torch.onnx.export(model, dummy_input, 'mnist.onnx', input_names=['input'], output_names=['output'])

print("✅ 模型已成功转换为 ONNX 格式！")

2. 构建 FastAPI 应用

from fastapi import FastAPI, File, UploadFile
from PIL import Image
import onnxruntime as ort
import numpy as np

app = FastAPI(title="手写数字识别 API")
ort_session = ort.InferenceSession('mnist.onnx')

@app.post("/predict")
async def predict(file: UploadFile = File(...)):
    # 读取上传的图片
    image = Image.open(file.file).convert('L')
    # 图片预处理
    image = image.resize((28, 28))
    image_data = np.array(image).astype(np.float32).reshape(1, 1, 28, 28)
    # 归一化
    image_data /= 255.0
    # 推理
    outputs = ort_session.run(None, {'input': image_data})
    prediction = np.argmax(outputs[0])
    # 返回结果
    return {"prediction": int(prediction)}

3. 测试

使用工具（如 cURL、Postman）发送请求，验证接口功能。

curl -X POST "http://localhost:8000/predict" -H "accept: application/json" -H "Content-Type: multipart/form-data" -F "file=@test_digit.png;type=image/png"

示例输出：

{
    "prediction": 7
}

小结

通过遵循上述最佳实践，我们可以简化部署流程，提高模型的推理性能，增强服务的可靠性和安全性。当然，在实际应用中，我们还需要根据具体情况进行优化和调整，希望本篇文章可以对各位读者有所帮助！