Web框架到高性能推理的完整AI应用部署流程——Flask / Django (Web框架)+ONNX (模型格式中间件)+TensorRT (高性能推理引擎)
·
一个常见的工作流是:用PyTorch训练模型 → 导出为ONNX格式 → 使用TensorRT优化并转换为TensorRT引擎 → 通过Flask/Django创建API来加载TensorRT引擎并处理推理请求
Flask / Django (Web框架):提供对外服务的接口和业务逻辑处理。用户通过浏览器或APP访问它提供的API,它负责接收数据、调用模型、返回结果。
- Flask: 是一个“微”框架。轻量、灵活、自由度高,核心功能简单,其他功能需要通过扩展来实现。非常适合快速构建小型项目或微服务API。
- Django:是一个“全栈”框架。功能强大、内置了ORM、管理员后台、用户认证、模板引擎等大量组件,开箱即用。适合构建中大型、复杂的Web应用。
from flask import Flask,request,jsonify
import your_model_module # 这里是你加载和推理模型的模块
app = Flask(__name__)
model = your_model_module.load_model() # 在启动时加载模型,避免每次请求都加载
@app.route('/predict',methods=['POST'])
def predict():
# 1. 接收请求数据
data = request.get_json()
input_data = data['input']
# 2. 调用模型进行推理
prediction = model.predict(input_data)
# 3. 返回结果
return jsonify({'prediction': prediction.tolist()})
客户端向 http://your-server-ip:5000/predict 发送一个POST请求,Body中包含JSON格式的输入数据,即可获得预测结果。
ONNX (模型格式中间件):作为模型转换的标准格式。它让不同框架(PyTorch, TensorFlow等)训练的模型可以有一个统一的格式,便于在不同硬件和推理引擎上运行。
核心使用方式是 导出 和 运行
import torch
import torchvision
# 加载一个预训练模型或你自己的模型
model = torchvision.models.resnet50(pretrained=True)
model.eval()
# 创建一个示例输入(dummy input)
dummy_input = torch.randn(1, 3, 224, 224)
# 导出模型为ONNX格式
torch.onnx.export(
model, # 要导出的模型
dummy_input, # 模型输入示例
"resnet50.onnx", # 导出的ONNX文件名
export_params=True, # 同时导出训练好的参数
opset_version=11, # ONNX算子集版本
do_constant_folding=True, # 是否执行常量折叠优化
input_names = ['input'], # 输入节点名称
output_names = ['output'] # 输出节点名称
)
import onnxruntime as ort
import numpy as np
# 创建ONNX Runtime推理会话
sess = ort.InferenceSession("resnet50.onnx")
# 准备输入数据(注意名称需与导出时定义的input_names一致)
input_name = sess.get_inputs()[0].name
input_data = np.random.randn(1, 3, 224, 224).astype(np.float32)
# 进行推理
outputs = sess.run(None, {input_name: input_data})
print(outputs[0])
TensorRT (高性能推理引擎):负责在NVIDIA GPU上极致优化和加速模型推理。它将ONNX或其它格式的模型转换成高度优化的引擎,实现最低延迟和最高吞吐量。
通过层融合(Layer Fusion)、精度校准(Precision Calibration(FP16/INT8))、内核自动调优(Kernel Auto-Tuning)等技术,大幅提升模型在NVIDIA GPU上的推理速度
通常有两种工作流:
- ONNX -> TensorRT(最常用):先将模型转换为ONNX,再用TensorRT的trtexec工具或Python API将ONNX模型构建和优化成TensorRT引擎(.engine文件)。
- 直接解析:TensorRT也可以直接解析PyTorch(通过Torch-TensorRT)或TensorFlow的模型。
# ONNX -> TensorRT Python API
import tensorrt as trt
# 1. 创建一个日志记录器和一个构建器
logger = trt.Logger(trt.Logger.WARNING)
builder = trt.Builder(logger)
# 2. 创建网络定义
network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
parser = trt.OnnxParser(network, logger)
# 3. 解析ONNX模型文件
with open(“your_model.onnx”, “rb”) as model:
if not parser.parse(model.read()):
for error in range(parser.num_errors):
print(parser.get_error(error))
# 4. 创建构建配置,设置精度等
config = builder.create_builder_config()
config.set_flag(trt.BuilderFlag.FP16) # 启用FP16精度加速
# 5. 构建并序列化引擎
serialized_engine = builder.build_serialized_network(network, config)
# 6. 保存引擎文件
with open(“model.engine”, “wb”) as f:
f.write(serialized_engine)
# --- 推理阶段 ---
# 7. 加载引擎文件并创建执行上下文
runtime = trt.Runtime(logger)
with open(“model.engine”, “rb”) as f:
engine = runtime.deserialize_cuda_engine(f.read())
context = engine.create_execution_context()
# 8. 为输入输出分配GPU内存,并进行推理(需配合PyCUDA或CuPy等库)
# ... (代码较为复杂,涉及GPU内存分配和数据拷贝)
更简单的方法:使用NVIDIA提供的命令行工具 trtexec 进行转换和性能测试。
trtexec --onnx=your_model.onnx --saveEngine=model.engine --fp16
更多推荐


所有评论(0)