ONNX Runtime与Flask/FastAPI：构建高性能AI推理API服务

ONNX Runtime与Flask/FastAPI：构建高性能AI推理API服务

【免费下载链接】onnxruntime microsoft/onnxruntime: 是一个用于运行各种机器学习模型的开源库。适合对机器学习和深度学习有兴趣的人，特别是在开发和部署机器学习模型时需要处理各种不同框架和算子的人。特点是支持多种机器学习框架和算子，包括 TensorFlow、PyTorch、Caffe 等，具有高性能和广泛的兼容性。 项目地址: https://gitcode.com/GitHub_Trending/on/onnxruntime

在AI应用开发中，将训练好的模型快速部署为API服务是落地的关键环节。ONNX Runtime（Open Neural Network Exchange Runtime）作为跨平台高性能推理引擎，与轻量级Web框架Flask/FastAPI结合，能构建出高效、易用的AI服务。本文将从环境搭建到性能优化，完整演示如何将ONNX模型转化为生产级API服务。

核心优势与应用场景

ONNX Runtime通过统一的ONNX格式支持TensorFlow、PyTorch等多框架模型，配合Flask/FastAPI的异步处理能力，特别适合以下场景：

• 实时推理服务（如人脸识别、文本分类）
• 边缘设备轻量化部署
• 多模型组合服务（如NLP+CV融合应用）

环境准备与基础依赖

核心组件安装

# 安装ONNX Runtime（CPU版本）
pip install onnxruntime
# 安装Web框架（二选一）
pip install flask  # 轻量级同步框架
# 或
pip install fastapi uvicorn  # 高性能异步框架

项目结构设计

ai-inference-service/
├── models/                 # 存放ONNX模型文件
│   └── resnet50.onnx
├── app.py                  # Flask/FastAPI应用
├── requirements.txt        # 依赖清单
└── utils/                  # 辅助函数（预处理/后处理）

快速实现：从模型到API

1. ONNX Runtime基础推理流程

ONNX Runtime的Python API提供简洁的模型加载与推理接口：

import onnxruntime as ort

# 创建推理会话（支持CPU/GPU）
sess_options = ort.SessionOptions()
sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
session = ort.InferenceSession(
    "models/resnet50.onnx",
    sess_options=sess_options,
    providers=["CPUExecutionProvider"]  # 可选"CUDAExecutionProvider"
)

# 准备输入数据（需匹配模型输入维度）
input_name = session.get_inputs()[0].name
input_data = np.random.randn(1, 3, 224, 224).astype(np.float32)

# 执行推理
outputs = session.run(None, {input_name: input_data})

2. Flask同步API实现

from flask import Flask, request, jsonify
import onnxruntime as ort
import numpy as np

app = Flask(__name__)
# 加载模型（全局单例避免重复初始化）
session = ort.InferenceSession("models/resnet50.onnx")
input_name = session.get_inputs()[0].name
output_name = session.get_outputs()[0].name

@app.route('/predict', methods=['POST'])
def predict():
    # 接收图像数据并预处理
    data = request.json['image']
    input_data = np.array(data, dtype=np.float32).reshape(1, 3, 224, 224)
    
    # 执行推理
    result = session.run([output_name], {input_name: input_data})
    
    # 返回预测结果
    return jsonify({"class_id": int(np.argmax(result[0]))})

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=5000)

3. FastAPI异步优化实现

FastAPI通过异步处理提升并发性能，适合高流量场景：

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

app = FastAPI()
session = ort.InferenceSession("models/resnet50.onnx")
input_name = session.get_inputs()[0].name
output_name = session.get_outputs()[0].name

class ImageRequest(BaseModel):
    image: list  # 接收图像像素列表

@app.post("/predict")
async def predict(request: ImageRequest):
    # 异步预处理（避免阻塞事件循环）
    loop = asyncio.get_event_loop()
    input_data = await loop.run_in_executor(
        None, 
        lambda: np.array(request.image, dtype=np.float32).reshape(1, 3, 224, 224)
    )
    
    # 执行推理（CPU推理为同步操作）
    result = session.run([output_name], {input_name: input_data})
    
    return {"class_id": int(np.argmax(result[0]))}

# 启动命令：uvicorn app:app --host 0.0.0.0 --port 8000

性能优化策略

1. 推理会话配置优化

# 启用图优化（默认开启）
sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
# 设置线程数（根据CPU核心数调整）
sess_options.intra_op_num_threads = 4
sess_options.inter_op_num_threads = 2
# 启用内存复用
sess_options.enable_cpu_mem_arena = True

2. 输入输出数据处理加速

• 使用numpy数组而非Python列表存储数据
• 对图像类任务采用OpenCV进行预处理（C++加速）
• 大模型考虑启用ONNX Runtime的IO绑定功能：

io_binding = session.io_binding()
io_binding.bind_cpu_input(input_name, input_data)
io_binding.bind_output(output_name, "cpu")
session.run_with_iobinding(io_binding)
result = io_binding.copy_outputs_to_cpu()[0]

3. 部署架构优化

对于高并发场景，推荐：

• FastAPI + Uvicorn（异步服务器）
• 多实例负载均衡（Nginx）
• GPU加速（安装onnxruntime-gpu并指定providers=["CUDAExecutionProvider"]）

常见问题与解决方案

Q: 如何处理模型加载耗时问题？

A: 采用应用启动时预加载模型，避免运行时动态加载：

# Flask应用中
def create_app():
    app = Flask(__name__)
    # 在应用创建时加载模型
    app.session = ort.InferenceSession("models/resnet50.onnx")
    return app

Q: 如何支持多模型并行服务？

A: 使用模型池或独立服务进程，避免单个会话占用过多资源。可参考官方Node.js示例的多模型管理模式：samples/nodejs/01_basic-usage/index.js

Q: GPU推理时报错"CUDA out of memory"怎么办？

A: 尝试：

1. 减少批处理大小
2. 启用混合精度推理
3. 使用模型量化工具减小模型体积：onnxruntime.quantize

部署与监控

Docker容器化

FROM python:3.9-slim
WORKDIR /app
COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt
COPY . .
CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "8000"]

性能监控

• 启用ONNX Runtime日志：

import onnxruntime as ort
ort.set_default_logger_severity(0)  # 0=VERBOSE, 1=INFO, 2=WARNING, 3=ERROR

• 集成Prometheus监控推理延迟和吞吐量

总结与扩展方向

本文介绍了基于ONNX Runtime和Flask/FastAPI构建AI推理API的完整流程，包括：

• 环境搭建与基础实现
• 同步/异步API设计
• 性能优化与部署策略

进阶探索方向：

• 模型热更新机制实现
• A/B测试框架集成
• 边缘设备部署（ONNX Runtime Mobile）

通过合理配置与优化，ONNX Runtime与现代Web框架的组合能够满足从原型验证到生产部署的全流程需求，为AI应用提供高效可靠的推理服务。

点赞+收藏+关注，获取更多AI部署实践技巧！下期预告：《ONNX模型量化与TensorRT加速实战》

【免费下载链接】onnxruntime microsoft/onnxruntime: 是一个用于运行各种机器学习模型的开源库。适合对机器学习和深度学习有兴趣的人，特别是在开发和部署机器学习模型时需要处理各种不同框架和算子的人。特点是支持多种机器学习框架和算子，包括 TensorFlow、PyTorch、Caffe 等，具有高性能和广泛的兼容性。 项目地址: https://gitcode.com/GitHub_Trending/on/onnxruntime