第二章:环境搭建与工具链

CUDA、cuDNN、TensorRT 的版本匹配

说实话,环境搭建这事儿,看着简单,坑是真不少。我见过太多人花了一周时间调不通模型,最后发现是版本没对上。

先说说核心原则:版本必须严格匹配。这不是闹着玩的。CUDA、cuDNN、TensorRT 三者之间,有明确的依赖关系。你想想看,TensorRT 底层调用的是 cuDNN 的算子,cuDNN 又依赖 CUDA 驱动。任何一个版本不对,都会导致推理失败。

⚠️ 我曾经踩过的坑: 有一次我图省事,直接装了最新版 CUDA 12.x,结果 TensorRT 8.6 根本不支持。折腾了两天,最后老老实实降级到 CUDA 11.8。嗯,从那以后我再也不敢乱来了。

推荐版本组合(截至 2024 年)

TensorRT 版本 CUDA 版本 cuDNN 版本 适用场景
8.6.x 11.8 8.7 生产环境最稳定
8.5.x 11.7 8.6 旧项目兼容
10.0.x 12.2 9.0 新特性尝鲜

我个人习惯用 TensorRT 8.6 + CUDA 11.8 + cuDNN 8.7 这套组合。为什么?因为它在稳定性和性能之间取得了很好的平衡。你如果做线上部署,我建议你也用这套。

安装步骤(以 Ubuntu 20.04 为例)

安装其实不复杂,但顺序很重要。我一般这么干:

  1. 先装 CUDA:去 NVIDIA 官网下载 runfile 或 deb 包。我个人喜欢用 runfile,因为可以指定安装路径,方便多版本共存。
  2. 再装 cuDNN:下载对应 CUDA 版本的 cuDNN 包,解压后把文件复制到 CUDA 目录下。
  3. 最后装 TensorRT:下载 tar 包,解压到指定目录,配置环境变量。

关键命令示例:

# 安装 CUDA 11.8
sudo sh cuda_11.8.0_520.61.05_linux.run

# 安装 cuDNN 8.7
tar -xzvf cudnn-linux-x86_64-8.7.0.84_cuda11-archive.tar.xz
sudo cp cudnn-*-archive/include/cudnn*.h /usr/local/cuda-11.8/include
sudo cp -P cudnn-*-archive/lib/libcudnn* /usr/local/cuda-11.8/lib64

# 安装 TensorRT 8.6
tar -xzvf TensorRT-8.6.1.6.Linux.x86_64-gnu.cuda-11.8.tar.gz
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/path/to/TensorRT-8.6.1.6/lib

💡 小技巧: 我习惯把 TensorRT 的 lib 路径加到 ~/.bashrc 里,这样每次打开终端就不用重新 export 了。省心。

验证安装是否成功

装完了怎么知道对不对?跑几个命令验证一下:

# 检查 CUDA
nvcc --version

# 检查 cuDNN
cat /usr/local/cuda/include/cudnn_version.h | grep CUDNN_MAJOR -A 2

# 检查 TensorRT
dpkg -l | grep TensorRT
# 或者
python3 -c "import tensorrt; print(tensorrt.__version__)"

如果都能正常输出版本号,恭喜你,环境搭好了。如果报错,别慌。我遇到过最多的问题是 LD_LIBRARY_PATH 没设对,或者 cuDNN 的 so 文件权限不对。检查一下这两点,基本能解决 80% 的问题。

trtexec 命令行工具入门

trtexec 是 TensorRT 自带的命令行工具。说白了,它就是帮你快速测试模型转换和推理性能的瑞士军刀。我平时做模型优化,第一步就是用 trtexec 跑个基准。

它的核心功能有三个:

  • 模型转换:把 ONNX 模型转成 TensorRT 的 engine 文件
  • 性能测试:跑推理,输出延迟和吞吐量
  • 精度验证:对比 FP32 和 FP16/INT8 的推理结果

常用命令示例

我直接给你几个我常用的命令,你拿去就能用:

# 1. 把 ONNX 转成 FP16 engine
trtexec --onnx=model.onnx --saveEngine=model_fp16.engine --fp16

# 2. 测试 engine 性能
trtexec --loadEngine=model_fp16.engine --shapes=input:1x3x224x224

# 3. 指定 batch size 测试
trtexec --loadEngine=model_fp16.engine --shapes=input:8x3x224x224

# 4. 查看模型层信息
trtexec --onnx=model.onnx --verbose --exportLayerInfo=layer_info.txt

⚠️ 注意: --shapes 参数必须跟模型输入的实际 shape 一致。我见过有人把 batch 维度写错了,结果推理结果全是错的。嗯,检查一下。

性能测试结果解读

跑完 trtexec 后,你会看到类似这样的输出:

[I] GPU Compute Time: min = 1.23 ms, max = 1.45 ms, mean = 1.31 ms
[I] Throughput: 763.45 qps

这里有两个关键指标:

  • GPU Compute Time:单次推理的 GPU 计算时间。这个值越小越好。
  • Throughput:每秒能处理的请求数。这个值越大越好。

我个人习惯看 mean 值,因为它能反映大多数情况下的性能。min 和 max 受系统波动影响大,参考价值有限。

💡 避坑指南: 我曾经在测试时发现 throughput 忽高忽低,后来发现是 GPU 频率没锁定。建议测试前先跑个 nvidia-smi -ac 5001,1590 锁定频率,这样结果才稳定。

trtexec 的高级用法

如果你觉得上面那些还不够,trtexec 还有一些隐藏功能:

  • 多流测试--streams=4 可以模拟多路并发推理
  • DLA 测试--useDLACore=0 可以测试 Jetson 设备上的 DLA 性能
  • INT8 校准--calib=calibration_data 可以指定校准数据集

这些功能我后面章节会详细讲,今天先混个脸熟。

总结一下

环境搭建这事儿,说白了就是「版本匹配 + 路径配置」。你只要记住:先查官方文档的版本兼容表,再动手装,基本不会出大问题。trtexec 是个好工具,多跑跑,你就能摸清模型的性能底细了。

下一章,我会带你手把手把 ONNX 模型转成 TensorRT engine,并对比不同精度的性能差异。到时候见。