必看！AI算法部署终极方案：PyTorch转ONNX+TensorRT加速，速度暴涨10倍+

renhongxia1

479人浏览 · 2026-05-28 09:03:48

renhongxia1 · 2026-05-28 09:03:48 发布

做AI算法部署的你，是不是常被这些问题逼到崩溃？PyTorch模型训练得再好，部署到实际场景就“掉链子”，推理速度慢到无法落地；转ONNX格式时频繁报错，要么维度不匹配，要么算子不支持；好不容易转成ONNX，对接TensorRT又卡壳，配置参数一头雾水，加速效果远不如预期；找遍教程不是残缺不全，就是版本过时，折腾几天还是原地踏步？

如果你也深陷“训练好模型却部署不了”的困境，别再独自踩坑！今天这篇PyTorch模型转ONNX+TensorRT加速全实操指南，就是为你量身打造的——不搞空洞理论，全程手把手带练，从环境配置、模型转换、参数优化到加速验证，每一步都有具体操作+代码实现+避坑要点，跟着做就能让你的AI模型推理速度暴涨，轻松落地工业级场景！

一、先搞懂：为什么必选ONNX+TensorRT组合做部署加速？

在AI算法部署领域，ONNX+TensorRT堪称“黄金组合”，优势无可替代：ONNX作为通用模型格式，能完美衔接PyTorch、TensorFlow等主流框架，解决不同框架模型的兼容性问题，是模型跨平台部署的“桥梁”；而TensorRT是NVIDIA专属的高性能推理引擎，能通过算子融合、量化感知训练、层间优化等核心技术，最大限度挖掘GPU算力，让模型推理速度实现质的飞跃，比原生PyTorch推理快5-10倍，甚至更高！

更关键的是，这套组合适配绝大多数工业级场景，无论是自动驾驶、智能安防，还是推荐系统、语音识别，只要用到PyTorch模型部署，都能通过它实现高效加速，是算法工程师必备的核心技能。

二、实操干货：PyTorch转ONNX+TensorRT加速全流程（附完整代码）

1. 环境准备：2步搞定配置，避开版本兼容坑

核心依赖：Python 3.8+、PyTorch 1.10+、ONNX 1.12+、TensorRT 8.4+（需匹配CUDA版本，建议CUDA 11.6+），推荐用Anaconda创建独立环境，避免依赖冲突。

bash
# 1. 创建并激活环境
conda create -n torch2trt python=3.9
conda activate torch2trt

# 2. 安装核心依赖（按CUDA 11.6版本适配）
pip install torch==1.13.1+cu116 torchvision==0.14.1+cu116 torchaudio==0.13.1 --extra-index-url https://download.pytorch.org/whl/cu116
pip install onnx==1.12.0 onnxruntime-gpu==1.14.1
# TensorRT需手动下载对应版本安装，安装后配置环境变量
pip install nvidia-pyindex
pip install nvidia-tensorrt==8.5.3.1

避坑要点：TensorRT版本必须与CUDA、PyTorch版本严格匹配，否则会出现“CUDA error”“算子不支持”等问题；若安装TensorRT失败，可直接从NVIDIA官网下载对应版本的tar包，解压后添加环境变量即可。

2. 第一步：PyTorch模型转ONNX格式（关键步骤+代码）

ONNX转换的核心是指定输入维度、动态维度适配（可选），这里以ResNet50模型为例，给出通用转换代码：

python
import torch
import torchvision.models as models

# 1. 加载预训练的PyTorch模型（或自定义模型）
model = models.resnet50(pretrained=True)
model.eval() # 切换到评估模式，避免BatchNorm等层影响转换

# 2. 构造虚拟输入（需与模型实际输入维度一致，通道顺序：RGB，格式：NCHW）
batch_size = 1
input_channel = 3
input_height = 224
input_width = 224
dummy_input = torch.randn(batch_size, input_channel, input_height, input_width)

# 3. 定义输出ONNX文件名
onnx_filename = "resnet50.onnx"

# 4. 执行转换（支持动态维度，这里指定batch_size为动态）
torch.onnx.export(
    model=model,
    args=dummy_input,
    f=onnx_filename,
    input_names=["input"], # 输入节点名称
    output_names=["output"], # 输出节点名称
    dynamic_axes={ # 动态维度配置，batch_size可动态变化
        "input": {0: "batch_size"},
        "output": {0: "batch_size"}
    },
    opset_version=12 # opset版本，建议11-13，过高可能不兼容TensorRT
)

# 5. 验证ONNX模型有效性
import onnx
onnx_model = onnx.load(onnx_filename)
onnx.checker.check_model(onnx_model)
print("PyTorch转ONNX成功！")

避坑要点：转换前必须将模型设为eval模式；虚拟输入维度需与实际业务场景一致；动态维度仅需指定必要的维度（如batch_size），过多动态维度会影响后续TensorRT加速效果。

3. 第二步：ONNX模型转TensorRT引擎（2种方式，按需选择）

TensorRT转换支持Python API和命令行两种方式，这里分别给出实操代码，新手推荐先试Python API，更易调试。

方式1：Python API转换（推荐，可自定义优化参数）

python
import tensorrt as trt

# 1. 初始化TensorRT logger
TRT_LOGGER = trt.Logger(trt.Logger.WARNING) # 只输出警告和错误信息
builder = trt.Builder(TRT_LOGGER)
network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
parser = trt.OnnxParser(network, TRT_LOGGER)

# 2. 解析ONNX模型
onnx_filename = "resnet50.onnx"
with open(onnx_filename, "rb") as model_file:
    parser.parse(model_file.read())

# 3. 配置TensorRT优化参数
config = builder.create_builder_config()
config.max_workspace_size = 1 << 30 # 最大工作空间（1GB），越大优化效果越好
# 开启FP16量化加速（精度损失小，速度提升明显）
if builder.platform_has_fast_fp16:
    config.set_flag(trt.BuilderFlag.FP16)

# 4. 构建TensorRT引擎并序列化保存
engine = builder.build_serialized_network(network, config)
trt_engine_filename = "resnet50_trt.engine"
with open(trt_engine_filename, "wb") as f:
    f.write(engine)

print("ONNX转TensorRT引擎成功！")

方式2：命令行转换（快速便捷，适合简单模型）

bash
# 基本命令（开启FP16加速，指定最大工作空间1GB）
trtexec --onnx=resnet50.onnx --saveEngine=resnet50_trt.engine --fp16 --workspace=1024

# 若需动态batch_size，添加以下参数（指定最小、最优、最大batch_size）
trtexec --onnx=resnet50.onnx --saveEngine=resnet50_trt.engine --fp16 --workspace=1024 --minShapes=input:1x3x224x224 --optShapes=input:8x3x224x224 --maxShapes=input:32x3x224x224

4. 第三步：TensorRT引擎推理验证（加速效果实测）

用转换后的TensorRT引擎推理，对比原生PyTorch的速度差异：

python
import torch
import tensorrt as trt
import pycuda.driver as cuda
import pycuda.autoinit
import numpy as np
import time

# 1. 加载TensorRT引擎
trt_engine_filename = "resnet50_trt.engine"
runtime = trt.Runtime(TRT_LOGGER)
with open(trt_engine_filename, "rb") as f:
    engine = runtime.deserialize_cuda_engine(f.read())
context = engine.create_execution_context()

# 2. 准备输入输出数据（GPU内存分配）
input_batch = torch.randn(8, 3, 224, 224).cuda() # batch_size=8
input_host = input_batch.cpu().numpy().astype(np.float32)
input_device = cuda.mem_alloc(input_host.nbytes)
output_host = np.empty((8, 1000), dtype=np.float32)
output_device = cuda.mem_alloc(output_host.nbytes)

# 3. TensorRT推理（多次运行取平均速度）
cuda.memcpy_htod(input_device, input_host)
context.set_binding_shape(0, (8, 3, 224, 224)) # 设置实际batch_size
start_time = time.time()
for _ in range(100):
    context.execute_v2([int(input_device), int(output_device)])
    cuda.memcpy_dtoh(output_host, output_device)
trt_time = (time.time() - start_time) / 100

# 4. 原生PyTorch推理速度对比
model = models.resnet50(pretrained=True).eval().cuda()
start_time = time.time()
for _ in range(100):
    with torch.no_grad():
        torch_output = model(input_batch)
torch_time = (time.time() - start_time) / 100

# 5. 输出速度对比结果
print(f"TensorRT平均推理时间：{trt_time*1000:.2f}ms")
print(f"原生PyTorch平均推理时间：{torch_time*1000:.2f}ms")
print(f"加速倍数：{torch_time/trt_time:.2f}倍")