📚 课程目录第2章环境准备

2、环境准备：安装TensorRT、验证trtexec可用性、获取示例模型（如ONNX、TensorFlow、PyTorch导出）的方法

TensorRT性能Benchmark的第一步，是搭建一个干净、可复现的软硬件环境。很多初学者在拿到NVIDIA GPU后，直接运行trtexec却报错“libnvinfer.so找不到”或“CUDA driver version is insufficient”，根源往往在于没有理清TensorRT、CUDA、cuDNN、显卡驱动之间的版本依赖关系。本章将围绕“安装TensorRT”“验证trtexec可用性”“获取示例模型”三个核心任务展开，帮助你从零构建一个可用于后续所有Benchmark实验的标准化环境。

上图展示了环境准备的三条主线：安装TensorRT、验证trtexec、获取示例模型。三者之间存在严格的先后依赖关系——只有TensorRT正确安装，trtexec才能运行；只有trtexec可用，才能对后续获取的模型进行Benchmark。下面逐一展开。

2.1 安装TensorRT：版本匹配是首要原则

核心原理：TensorRT是NVIDIA推出的高性能深度学习推理SDK，它依赖于CUDA和cuDNN，而CUDA又依赖于显卡驱动。因此，安装TensorRT的本质是确保“驱动 → CUDA → cuDNN → TensorRT”这条链路版本兼容。NVIDIA官方提供了兼容性矩阵，但更直接的方法是使用nvidia-smi查看驱动支持的CUDA版本，再据此选择TensorRT版本。

安装方式	适用场景	优点	缺点
Debian/RPM包	Ubuntu/CentOS服务器	自动处理依赖，升级方便	需要root权限，版本固定
Tar压缩包	任意Linux环境	解压即用，可自定义路径	需手动配置环境变量
Python pip包	快速原型开发	安装简单，与Python生态集成	仅包含Python API，trtexec需单独获取
Docker镜像	团队协作/CI/CD	环境隔离，版本可控	镜像体积大，GPU透传需额外配置

操作流程（以Ubuntu 20.04 + TensorRT 8.6为例）：

# 1. 检查驱动和CUDA
nvidia-smi  # 查看驱动版本和CUDA版本
nvcc --version  # 查看CUDA Toolkit版本

# 2. 下载TensorRT Tar包（需NVIDIA账号）
# 从 https://developer.nvidia.com/tensorrt 下载对应CUDA版本的包
tar -xzvf TensorRT-8.6.1.6.Ubuntu-20.04.x86_64-gnu.cuda-11.8.tar.gz

# 3. 配置环境变量（写入~/.bashrc）
export TRT_ROOT=/path/to/TensorRT-8.6.1.6
export LD_LIBRARY_PATH=$TRT_ROOT/lib:$LD_LIBRARY_PATH
export PATH=$TRT_ROOT/bin:$PATH

# 4. 验证安装
python3 -c "import tensorrt; print(tensorrt.__version__)"  # 应输出8.6.1.6

常见误区：不要只看TensorRT版本号，必须同时确认CUDA和cuDNN版本。例如TensorRT 8.6要求CUDA 11.8和cuDNN 8.9，如果系统只有CUDA 11.4，即使强行安装也会在运行时崩溃。建议使用ldd命令检查依赖库是否齐全。

2.2 验证trtexec可用性：从--help到首次Benchmark

核心原理：trtexec是TensorRT自带的命令行工具，用于构建引擎、执行推理和输出性能指标。它位于$TRT_ROOT/bin目录下。验证trtexec可用性不仅是确认它能运行，更重要的是确认它能正确识别GPU、加载ONNX模型并生成可执行的引擎。

验证步骤：

基础验证：运行trtexec --help，检查参数列表是否完整。如果输出为空或报错，说明环境变量配置有误。
版本确认：运行trtexec --version，确认版本号与安装的TensorRT一致。
硬件识别：运行trtexec --listDevices，确认能列出所有GPU设备。
端到端验证：使用一个简单的ONNX模型（如ResNet-50）执行trtexec --onnx=model.onnx --saveEngine=model.engine，观察是否成功生成引擎文件。

# 端到端验证命令
trtexec --onnx=resnet50.onnx --saveEngine=resnet50.engine --fp16

# 预期输出片段
[I] Engine built in 12.3456 sec.
[I] === Performance summary ===
[I] Throughput: 1250.67 qps
[I] Latency: min = 0.789 ms, max = 1.234 ms, mean = 0.912 ms

经验提醒：如果trtexec在构建引擎时卡住或报错，优先检查GPU显存是否充足（nvidia-smi查看）、模型输入shape是否与ONNX一致、是否缺少自定义插件。建议先用--verbose参数运行，获取详细日志。

2.3 获取示例模型：PyTorch/TensorFlow导出ONNX的完整方法

核心原理：TensorRT本身不直接支持PyTorch或TensorFlow模型，需要通过ONNX作为中间格式。ONNX（Open Neural Network Exchange）是一种开放的模型表示格式，能够保留计算图结构和权重。导出ONNX的关键在于：固定输入shape、指定动态维度、确保算子兼容。

2.3.1 PyTorch导出ONNX

import torch
import torchvision.models as models

# 加载模型并设为eval模式
model = models.resnet50(pretrained=True).eval()

# 构造示例输入（batch=1, 3通道, 224x224）
dummy_input = torch.randn(1, 3, 224, 224)

# 导出ONNX
torch.onnx.export(
    model,
    dummy_input,
    "resnet50.onnx",
    input_names=["input"],
    output_names=["output"],
    dynamic_axes={"input": {0: "batch"}, "output": {0: "batch"}},  # 动态batch
    opset_version=17
)
print("ONNX导出成功")

关键参数说明：

dynamic_axes：指定哪些维度是动态的。对于Benchmark，建议先固定batch=1，后续再测试动态batch。
opset_version：ONNX算子集版本，TensorRT 8.6支持opset 17及以下。版本过高可能导致算子不支持。
eval()模式：必须调用，否则BatchNorm等层会使用训练时的统计量。

2.3.2 TensorFlow导出ONNX

# 使用tf2onnx工具
import tensorflow as tf
import tf2onnx

# 加载SavedModel或Keras模型
model = tf.keras.applications.ResNet50(weights="imagenet")

# 转换为ONNX
spec = (tf.TensorSpec((None, 224, 224, 3), tf.float32, name="input"),)
onnx_model = tf2onnx.convert.from_keras(model, input_signature=spec, opset=17)
onnx_model[0].save("resnet50_tf.onnx")

  案例拆解：某团队需要将PyTorch训练的YOLOv5模型部署到TensorRT。初始导出时未设置dynamic_axes，导致模型只能处理固定尺寸的输入。后续修改为动态batch和动态宽高后，trtexec构建引擎时需指定--minShapes、--optShapes、--maxShapes参数，最终在A100上实现了2.3ms的端到端延迟，相比原始PyTorch推理提升了4.7倍。复盘发现，关键成功因素在于：①导出前验证了ONNX模型的输入输出节点名；②使用onnxsim简化了计算图；③在TensorRT中启用了FP16精度。

2.3.3 模型验证清单

检查项	工具/命令	通过标准
ONNX结构完整性	onnx.checker.check_model	无异常抛出
输入输出节点名	onnx.load + print(graph.input/output)	名称清晰，与代码一致
算子兼容性	trtexec --onnx=model.onnx --verbose	无“Unsupported Operator”错误
数值一致性	trtexec --loadEngine=model.engine --infer	输出与PyTorch/TF结果误差<1e-3

风险提示：ONNX导出后务必用onnxruntime或onnx.checker验证模型可运行。很多问题（如动态shape未正确声明、自定义算子缺失）在导出时不会报错，但会在TensorRT构建引擎时暴露。建议在导出后立即执行一次onnxruntime推理，确认输出与原始框架一致。

2.4 环境准备常见误区与总结复盘

误区一：认为TensorRT安装成功=trtexec可用。实际上，trtexec依赖libnvinfer.so和libnvonnxparser.so，如果LD_LIBRARY_PATH未正确设置，trtexec会报“cannot open shared object file”。

误区二：忽略CUDA版本匹配。TensorRT 8.x对CUDA版本有严格要求，例如8.6要求CUDA 11.8，8.5要求CUDA 11.7。使用nvcc --version确认CUDA Toolkit版本，而不是只看驱动版本。

误区三：导出ONNX时未固定输入shape。如果模型包含Reshape、Gather等对shape敏感的算子，动态shape可能导致TensorRT构建失败或性能下降。建议先固定shape完成Benchmark，再逐步引入动态维度。