第二章:环境搭建与工具链
CUDA、cuDNN、TensorRT 的版本匹配
说实话,环境搭建这事儿,看着简单,坑是真不少。我见过太多人花了一周时间调不通模型,最后发现是版本没对上。
先说说核心原则:版本必须严格匹配。这不是闹着玩的。CUDA、cuDNN、TensorRT 三者之间,有明确的依赖关系。你想想看,TensorRT 底层调用的是 cuDNN 的算子,cuDNN 又依赖 CUDA 驱动。任何一个版本不对,都会导致推理失败。
推荐版本组合(截至 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 为例)
安装其实不复杂,但顺序很重要。我一般这么干:
- 先装 CUDA:去 NVIDIA 官网下载 runfile 或 deb 包。我个人喜欢用 runfile,因为可以指定安装路径,方便多版本共存。
- 再装 cuDNN:下载对应 CUDA 版本的 cuDNN 包,解压后把文件复制到 CUDA 目录下。
- 最后装 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,并对比不同精度的性能差异。到时候见。