一个常见的工作流是:用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

更多推荐