TensorRT 学习文档（一）：从 PyTorch 训练到 TensorRT 部署的完整流程记录-平芜编程栈

TensorRT 学习文档（一）

从 PyTorch 训练到 TensorRT 部署的完整流程记录

1. 整体流程

PyTorch 模型 ↓ torch.onnx.export ONNX 模型 ↓ onnxsim（可选） 简化 ONNX ↓ trtexec / TensorRT API TensorRT Engine ↓ 部署 推理服务

每个阶段的产物：

阶段	产物	特点
PyTorch	`.pth`权重	训练用，Python 运行
ONNX	`.onnx`	通用计算图，跨框架
TensorRT	`.engine`	GPU 专用优化，硬件相关

2. 环境准备

2.1 版本要求

组件	版本	说明
Ubuntu	22.04	系统
CUDA	13.0	驱动 580+
TensorRT	11.4.0.106	通过 deb 本地仓库安装
CMake	3.31+	编译 sample 需要
Python	3.10	conda 环境

2.2 安装 TensorRT（deb 本地仓库）

# 1. 安装本地仓库包sudodpkg-inv-tensorrt-local-repo-ubuntu2204-10.16.1-cuda-13.2_1.0-1_amd64.deb# 2. 复制密钥sudocp/var/nv-tensorrt-local-repo-ubuntu2204-10.16.1-cuda-13.2/nv-tensorrt-local-F07E5E88-keyring.gpg /usr/share/keyrings/# 3. 更新软件列表sudoapt-getupdate# 4. 安装 TensorRTsudoapt-getinstalltensorrt# 5. 安装 Python 包（可选）sudoapt-getinstallpython3-libnvinfer

2.3 验证

# 命令行工具trtexec--help# C++ 库dpkg-l|greplibnvinfer# Python 包python3-c"import tensorrt; print(tensorrt.__version__)"

2.4 conda 环境安装（推荐）

conda create-ntrt_envpython=3.10-yconda activate trt_env pipinstalltensorrt-cu13 pycuda numpy onnx onnxruntime onnxsim polygraphy

3. PyTorch 模型基础

3.1 卷积层参数

nn.Conv2d(输入通道,输出通道,核大小)

参数量公式：

输出通道 × (输入通道 × 核高 × 核宽 + 1)

示例：

nn.Conv2d(16,32,3)# 参数量 = 32 × (16×3×3 + 1) = 4640

3.2 全连接层参数

nn.Linear(输入数,输出数)

参数量公式：

输出数 × (输入数 + 1)

示例：

nn.Linear(1568,10)# 参数量 = 10 × (1568 + 1) = 15690

3.3 数据流示例

classSimpleCNN(nn.Module):def__init__(self):super().__init__()self.conv1=nn.Conv2d(1,16,3,padding=1)self.conv2=nn.Conv2d(16,32,3,padding=1)self.relu=nn.ReLU()self.pool=nn.MaxPool2d(2,2)self.fc=nn.Linear(32*7*7,10)defforward(self,x):x=self.pool(self.relu(self.conv1(x)))# 16 × 14 × 14x=self.pool(self.relu(self.conv2(x)))# 32 × 7 × 7x=torch.flatten(x,1)# 1568x=self.fc(x)# 10returnx

数据流：

输入 1×28×28 ↓ conv1（1→16 层） 16×28×28 ↓ pool 16×14×14 ↓ conv2（16→32 层） 32×14×14 ↓ pool 32×7×7 ↓ flatten 1568 ↓ fc 10

3.4 卷积的权值共享

维度	是否共享
同一输出层，同一输入层，不同空间位置	✅ 共享
同一输出层，不同输入层	❌ 不共享
不同输出层，任意	❌ 不共享

滑动 = 同一套参数在不同位置重复使用。

4. PyTorch → ONNX

4.1 导出脚本

importtorchfrommodel_v2importSimpleCNN# 1. 准备模型model=SimpleCNN()model.eval()# 必须！切换到推理模式# 2. 准备假输入dummy_input=torch.randn(1,1,28,28)# 3. 导出torch.onnx.export(model,dummy_input,"model_v2.onnx",input_names=["input"],output_names=["output"],opset_version=17)

4.2 关键参数

参数	作用
`model`	要导出的模型
`dummy_input`	假输入，决定图结构
`input_names`	给输入张量起名字
`output_names`	给输出张量起名字
`opset_version`	ONNX 操作集版本
`dynamic_axes`	指定动态维度

4.3 trace vs script

trace	script
原理	跑一遍，记录实际操作	读代码，解析成图
`if`	只记录走的分支	两条都保留
`for`	展开成具体次数	保留循环结构
`getattr`	能处理	❌ 报错
动态列表	能处理	❌ 报错
Python 特性	全部支持	只支持子集

选择原则：

纯顺序模型 → trace
有控制流 → script
有动态特性 → trace

trace 的风险：

ifx.sum()>0:x=self.fc1(x)else:x=self.fc2(x)

假输入走 fc1，导出后只有 fc1，fc2 完全丢失。静默错误，不报错。

4.4 打印 ONNX 图结构

importonnx model=onnx.load("model_v2.onnx")print("=== 输入 ===")forinpinmodel.graph.input:print(f"{inp.name}, shape={[d.dim_valuefordininp.type.tensor_type.shape.dim]}")print("\n=== 输出 ===")foroutinmodel.graph.output:print(f"{out.name}, shape={[d.dim_valuefordinout.type.tensor_type.shape.dim]}")print("\n=== 节点 ===")fori,nodeinenumerate(model.graph.node):print(f" [{i}]{node.op_type}输入:{list(node.input)}输出:{list(node.output)}")

5. ONNX → TensorRT

5.1 基本转换

trtexec--onnx=model_v2.onnx--saveEngine=model_v2.engine

5.2 转换过程做了什么

1. 解析 ONNX → 读出计算图 2. 构建网络 → 把 ONNX 节点转成 TensorRT 的层 3. 优化 → 选最快的算法，融合能融合的层 4. 生成 engine → 序列化保存

关键优化：

优化	说明
算子融合	Conv+BN+Relu 合成一个 kernel
算法选择	同一卷积有多种实现，选最快的
精度校准	FP32/FP16/INT8
内存优化	复用显存

5.3 层融合

能融合的层：

Relu、Sigmoid、BatchNorm（推理时）
Scale、Add（相同形状）
逐元素操作

不能融合的层：

Softmax（需要全局求和）
Reshape（改变数据布局）
Concat（需要等所有输入）

融合的条件：每个输出元素只依赖对应位置的输入，不需要等所有数据。

5.4 参数说明

参数	作用
`--onnx=`	输入 ONNX 文件
`--saveEngine=`	输出 engine 文件
`--minShapes/--optShapes/--maxShapes`	动态 shape 范围
`--fp16`	开启 FP16
`--int8`	开启 INT8 量化

5.5 动态 Shape 转换

trtexec--onnx=model_v2_dynamic.onnx\--saveEngine=model_v2_dynamic.engine\--minShapes=input:1x1x28x28\--optShapes=input:4x1x28x28\--maxShapes=input:16x1x28x28

minShapes：最小输入
optShapes：最常用输入（优化重点）
maxShapes：最大输入（显存按这个分配）

6. 推理运行时

6.1 核心概念

概念	说明
Engine	编译好的模型，只读，可序列化
Context	运行实例，有独立中间激活值
Binding	输入输出张量，需要自己分配内存
Stream	CUDA 执行队列，异步

6.2 关系图

Engine ├── Context 1 → Binding 地址 1 → Stream 1 ├── Context 2 → Binding 地址 2 → Stream 2 └── ...

Binding 定义公用（名字、形状），地址每个 context 独立。

6.3 完整推理示例

importtensorrtastrtimportpycuda.driverascudaimportpycuda.autoinitimportnumpyasnp# 1. 加载 enginelogger=trt.Logger(trt.Logger.WARNING)runtime=trt.Runtime(logger)withopen("model.engine","rb")asf:engine=runtime.deserialize_cuda_engine(f.read())# 2. 创建 contextcontext=engine.create_execution_context()# 3. 查输入输出input_names=[]output_names=[]foriinrange(engine.num_io_tensors):name=engine.get_tensor_name(i)mode=engine.get_tensor_mode(name)ifmode==trt.TensorIOMode.INPUT:input_names.append(name)else:output_names.append(name)# 4. 分配内存bufs={}fornameininput_names+output_names:shape=tuple(engine.get_tensor_shape(name))dtype=trt.nptype(engine.get_tensor_dtype(name))host=np.random.randn(*shape).astype(dtype)ifnameininput_names \elsenp.empty(shape,dtype=dtype)device=cuda.mem_alloc(host.nbytes)context.set_tensor_address(name,int(device))bufs[name]=(host,device)# 5. 执行stream=cuda.Stream()fornameininput_names:host,device=bufs[name]cuda.memcpy_htod_async(device,host,stream)context.execute_async_v3(stream_handle=stream.handle)fornameinoutput_names:host,device=bufs[name]cuda.memcpy_dtoh_async(host,device,stream)stream.synchronize()

6.4 单 context vs 多 context

配置	能否并行	原因
1 context + 1 stream	❌	只有一个工作台
1 context + 多 stream	❌ 不安全	中间激活值只有一份
多 context + 多 stream	✅	每个 context 独立内存

7. ONNX 图优化

7.1 三类优化

类型	说明
算子融合	Conv+BN+Relu → 一个节点
常量折叠	能在导出时算出来的，直接替换成常量
Shape 推理	静态推导每个节点的输出形状

7.2 用 onnxsim 简化

pipinstallonnxsim onnxsim model_v2.onnx model_v2_sim.onnx

输出示例：

┏━━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━┓ ┃ ┃ Original Model ┃ Simplified Model ┃ ┡━━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━┩ │ Concat │ 10 │ 5 │ │ Constant │ 91 │ 88 │ │ Conv │ 55 │ 55 │ └──────────────────┴─────────────────┴──────────────────┘

7.3 移除 initializer 中的输入

python3-monnxruntime.tools.remove_initializer_from_input\--inputvpd_d.onnx\--outputvpd_d_fixed.onnx

为什么需要：旧版 PyTorch 导出的模型，权重被同时声明为输入，影响优化。

7.4 什么时候需要 ONNX 优化

情况	需要吗
模型简单，TensorRT 一次转成功	❌ 可以跳过
模型复杂，TensorRT 报错	✅ 必须
想减少编译时间	建议
想减少运行时 kernel 数量	建议

8. 精度验证

8.1 为什么需要

TensorRT 转换可能引入误差：

层融合改变计算顺序
算法选择和原始实现不同
FP16/INT8 量化有精度损失

8.2 用 polygraphy 验证

polygraphy run model_v2.onnx\--trt\--onnxrt\--input-shapes"input:[1,1,28,28]"

预期输出：

[I] Absolute Difference | Stats: mean=2.1e-08, max=5.9e-08 [I] Relative Difference | Stats: mean=2.0e-07, max=5.5e-07 [I] PASSED | Output: 'output' | Difference is within tolerance [I] PASSED | Runtime: 5.8s

8.3 误差标准

精度	可接受误差
FP32	< 1e-5
FP16	< 1e-3
INT8	< 1e-2

8.4 查看模型信息

polygraphy inspect model model_v2.onnx

输出：

---- 1 Graph Input(s) ---- {data [dtype=float32, shape=(1, 3, 192, 576)]} ---- 5 Graph Output(s) ---- {p3_conv_conf [dtype=float32, shape=(1, 8, 48, 144)], p4_conv_conf [dtype=float32, shape=(1, 8, 24, 72)], ...}

8.5 polygraphy 与 engine 的关系

polygraphy 用--trt时，会自动从 ONNX 构建 TensorRT engine，不需要提前用trtexec转出.engine文件。

你跑的命令：

polygraphy run model_v2.onnx--trt--onnxrt--input-shapes"input:[1,1,28,28]"

它内部做了：

model_v2.onnx ↓ TensorRT ONNX Parser ↓ TensorRT Builder 临时 engine（在内存里，不保存） ↓ 跑推理，和 ONNX Runtime 对比

两种运行模式

模式 1：从 ONNX 现场构建（默认）

polygraphy run model.onnx--trt--onnxrt

不需要预先有 engine
每次跑都会重新构建，较慢
适合临时验证

模式 2：加载已有 engine

polygraphy run model.onnx--trt--onnxrt--load-engine=model.engine

或者直接对 engine 跑：

polygraphy run model.engine--trt

不需要重新构建，快
适合反复测试同一个 engine

是否自动保存 engine

默认不保存。临时构建的 engine 只存在内存里，跑完就释放，不会落盘。

想保存，加参数：

polygraphy run model.onnx\--trt\--onnxrt\--save-engine=model.engine

跑完当前目录会生成model.engine。下次可以：

polygraphy run model.onnx--trt--load-engine=model.engine

直接加载，不用重新构建，快很多。

为什么默认不保存

polygraphy 的定位是验证工具，不是转换工具。

它的主要任务是：

ONNX 输出 vs TensorRT 输出 → 对比 → PASSED / FAILED

engine 只是中间产物，验证完就丢。

要保存 engine，用trtexec --saveEngine=或加--save-engine。

对比

工具	默认保存 engine 吗
`trtexec --saveEngine=`	✅ 明确要求保存
`polygraphy run --trt`	❌ 不保存
`polygraphy run --trt --save-engine=`	✅ 保存

9. 动态 Shape

9.1 导出动态 ONNX

torch.onnx.export(model,dummy_input,"model_v2_dynamic.onnx",input_names=["input"],output_names=["output"],dynamic_axes={"input":{0:"batch_size"},"output":{0:"batch_size"}},opset_version=17)

9.2 转换时指定范围

trtexec--onnx=model_v2_dynamic.onnx\--saveEngine=model_v2_dynamic.engine\--minShapes=input:1x1x28x28\--optShapes=input:4x1x28x28\--maxShapes=input:16x1x28x28

9.3 验证不同 batch

trtexec--loadEngine=model_v2_dynamic.engine--shapes=input:1x1x28x28 trtexec--loadEngine=model_v2_dynamic.engine--shapes=input:8x1x28x28 trtexec--loadEngine=model_v2_dynamic.engine--shapes=input:16x1x28x28

9.4 动态 Shape 注意事项

maxShapes决定显存分配，设太大会浪费显存
optShapes是优化重点，设成最常用的 batch
超出maxShapes的输入会报错

10. 常见问题

10.1 安装问题

问题	原因	解决
`trtexec: command not found`	未安装或不在 PATH	`export PATH=/usr/src/tensorrt/bin:$PATH`
`libnvinfer.so: cannot open`	LD_LIBRARY_PATH 未设置	`export LD_LIBRARY_PATH=/usr/lib/x86_64-linux-gnu:$LD_LIBRARY_PATH`
Python 包找不到	conda 环境隔离	在 conda 环境里`pip install tensorrt-cu13`

10.2 转换问题

问题	原因	解决
`Unknown option: --explicitBatch`	TensorRT 11.x 已移除	去掉这个参数
`Unsupported operator`	算子不支持	用 onnxsim 简化，或自定义插件
`CMake 3.31 or higher is required`	CMake 版本低	升级 CMake

10.3 精度问题

问题	原因	解决
输出差异 > 1e-2	转换有问题	检查输入对齐、模型加载
某个输出差异大	该分支有算子问题	用 polygraphy 定位到具体层
Sigmoid 后差异大	正常	Sigmoid 会放大相对误差

10.4 性能问题

问题	原因	解决
GPU compute time unstable	模型太小，测量噪声	锁定 GPU 频率`sudo nvidia-smi -lgc 1500`
推理比预期慢	没开 FP16	加`--fp16`
显存不够	maxShapes 设太大	减小 maxShapes

附录：常用命令速查

# 导出 ONNXpython3 export_onnx.py# 简化 ONNXonnxsim model.onnx model_sim.onnx# 查看 ONNX 信息polygraphy inspect model model.onnx# 转换 TensorRTtrtexec--onnx=model.onnx--saveEngine=model.engine# 动态 Shape 转换trtexec--onnx=model.onnx--saveEngine=model.engine\--minShapes=input:1x1x28x28\--optShapes=input:4x1x28x28\--maxShapes=input:16x1x28x28# 精度验证（不保存 engine）polygraphy run model.onnx--trt--onnxrt--input-shapes"input:[1,1,28,28]"# 精度验证 + 保存 enginepolygraphy run model.onnx--trt--onnxrt--save-engine=model.engine --input-shapes"input:[1,1,28,28]"# 加载已有 engine 验证polygraphy run model.onnx--trt--load-engine=model.engine# 加载 engine 跑推理trtexec--loadEngine=model.engine# 指定 batch 跑推理trtexec--loadEngine=model.engine--shapes=input:8x1x28x28

附录：工具读法

工具	读法	音标
TensorRT	Tensor-R-T	/ˈtensər ɑːr tiː/
ONNX	欧尼克斯	/ˈɒnɪks/
trtexec	T-R-T-exec	/tiː ɑːr tiː ˈeksək/
polygraphy	波利格拉菲	/pəˈlɪɡrəfi/
onnxsim	ONNX-sim	/ˈɒnɪks sɪm/
PyCUDA	Py-CUDA	/paɪ ˈkuːdə/