2、环境搭建与工具链:CUDA、cuDNN、TensorRT的版本匹配,安装与验证,trtexec命令行工具初探

说实话,做TensorRT开发最让人头疼的,往往不是模型本身,而是环境搭建。我见过太多同事花了一整天,就为了把CUDA和TensorRT的版本对上号。嗯,这节课我们就来把这套东西彻底理清楚。

2.1 版本匹配:别让版本坑了你

先说说版本匹配这件事。我个人习惯是,每次装环境前先画个版本对应表。为什么?因为CUDA、cuDNN、TensorRT这三个东西,版本不对应的话,跑起来全是玄学错误。

举个例子,TensorRT 8.5 GA Update 2 这个版本,它要求CUDA 11.0到11.8之间。你装个CUDA 12.0试试?编译都过不去。我曾经在项目里吃过这个亏,花了两天排查,最后发现就是CUDA版本高了。

下面这张表是我整理的常用版本对应关系,你直接照着来就行:

TensorRT版本 CUDA版本要求 cuDNN版本要求
8.4 GA 11.0 - 11.7 8.4+
8.5 GA 11.0 - 11.8 8.6+
8.6 EA 11.6 - 12.0 8.7+
10.0 GA 11.8 - 12.2 8.9+
注意:千万别只看大版本号。比如CUDA 11.7和11.8,虽然只差一个小版本,但TensorRT 8.4就不支持11.8。我建议你直接去NVIDIA官网查官方兼容性矩阵,那个最准。

2.2 安装流程:一步步来,别跳步

安装顺序其实有讲究。我建议你按这个顺序来:CUDA → cuDNN → TensorRT。为什么?因为后面的依赖前面的。

2.2.1 安装CUDA

CUDA的安装其实挺简单的。去NVIDIA官网下载runfile或者deb包都行。我个人习惯用runfile,因为可以指定安装路径,方便管理多个版本。

# 下载CUDA 11.8 runfile
wget https://developer.download.nvidia.com/compute/cuda/11.8.0/local_installers/cuda_11.8.0_520.61.05_linux.run

# 执行安装
sudo sh cuda_11.8.0_520.61.05_linux.run

# 配置环境变量
export PATH=/usr/local/cuda-11.8/bin:$PATH
export LD_LIBRARY_PATH=/usr/local/cuda-11.8/lib64:$LD_LIBRARY_PATH
小技巧:安装时记得不要勾选驱动选项,除非你的显卡驱动版本太老。我遇到过有人装了CUDA自带的驱动,结果把原来的驱动覆盖了,显卡直接黑屏。

2.2.2 安装cuDNN

cuDNN其实就是一堆库文件。下载后解压,然后把文件复制到CUDA目录下就行。

# 解压cuDNN
tar -xzvf cudnn-linux-x86_64-8.6.0.163_cuda11-archive.tar.xz

# 复制到CUDA目录
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

# 设置权限
sudo chmod a+r /usr/local/cuda-11.8/include/cudnn*.h
sudo chmod a+r /usr/local/cuda-11.8/lib64/libcudnn*

验证一下是否装好了:

# 查看cuDNN版本
cat /usr/local/cuda-11.8/include/cudnn_version.h | grep CUDNN_MAJOR -A 2

2.2.3 安装TensorRT

TensorRT的安装稍微复杂一点。我建议你用tar包方式安装,这样方便切换版本。

# 解压TensorRT
tar -xzvf TensorRT-8.5.3.1.Linux.x86_64-gnu.cuda-11.8.cudnn8.6.tar.gz

# 配置环境变量
export TRT_ROOT=/path/to/TensorRT-8.5.3.1
export LD_LIBRARY_PATH=$TRT_ROOT/lib:$LD_LIBRARY_PATH
export PATH=$TRT_ROOT/bin:$PATH
重点:记得把TensorRT的Python包也装上,后面写推理代码要用。
pip install $TRT_ROOT/python/tensorrt-8.5.3.1-cp38-none-linux_x86_64.whl

2.3 验证环境:装没装对,跑一下就知道

环境装完了,怎么验证?别光看版本号,得真跑起来才算数。

2.3.1 验证CUDA

# 查看CUDA版本
nvcc --version

# 跑个简单的测试
nvidia-smi

如果nvidia-smi能正常显示GPU信息,说明驱动和CUDA基本没问题。

2.3.2 验证cuDNN

cuDNN的验证稍微麻烦点。我一般用它的sample程序来测:

# 编译cuDNN sample
cd /usr/local/cuda-11.8/samples
sudo cp -r /path/to/cudnn-samples/* .

# 编译并运行
make
./mnistCUDNN

如果输出显示"Test passed!",那cuDNN就装对了。

2.3.3 验证TensorRT

TensorRT的验证最简单,直接用Python:

import tensorrt as trt
print(trt.__version__)

# 创建builder看看能不能正常工作
builder = trt.Builder(trt.Logger(trt.Logger.WARNING))
print("TensorRT builder created successfully!")
注意:如果导入tensorrt时报错找不到libnvinfer.so,八成是LD_LIBRARY_PATH没配好。我曾经在这个问题上卡了半小时,最后发现是路径写错了。

2.4 trtexec命令行工具初探

trtexec是TensorRT自带的一个命令行工具,说白了就是个瑞士军刀。它能帮你做模型转换、性能测试、精度验证,甚至还能生成序列化引擎。

我第一次用trtexec的时候,觉得这东西太方便了。不用写一行代码,就能把ONNX模型转成TensorRT引擎。

2.4.1 基本用法

# 查看帮助
trtexec --help

# 最简单的用法:将ONNX模型转为TensorRT引擎
trtexec --onnx=model.onnx --saveEngine=model.engine

# 指定精度
trtexec --onnx=model.onnx --fp16 --saveEngine=model_fp16.engine

# 指定batch size
trtexec --onnx=model.onnx --minShapes=input:1x3x224x224 \
        --optShapes=input:8x3x224x224 \
        --maxShapes=input:16x3x224x224

2.4.2 性能测试

trtexec最强大的功能之一就是性能测试。它能告诉你模型的推理延迟和吞吐量。

# 跑100次推理,统计性能
trtexec --loadEngine=model.engine --iterations=100 --avgRuns=50

# 输出示例
# [I] Average on 50 runs - GPU latency: 2.34567 ms - Host latency: 2.45678 ms
# [I] Throughput: 426.78 qps
经验之谈:我一般用trtexec做性能基线测试。比如优化前延迟是5ms,优化后变成2ms,那效果就很明显了。但要注意,trtexec测的是纯推理时间,实际部署时还要考虑前后处理。

2.4.3 精度验证

用trtexec还能验证FP16或INT8量化后的精度损失:

# 对比FP32和FP16的精度
trtexec --loadEngine=model_fp32.engine --dumpOutput --exportOutput=fp32_output.json
trtexec --loadEngine=model_fp16.engine --dumpOutput --exportOutput=fp16_output.json

# 然后自己写个脚本对比两个json文件的差异

2.4.4 常用参数速查

参数 说明 示例
--onnx 指定ONNX模型路径 --onnx=model.onnx
--saveEngine 保存序列化引擎 --saveEngine=model.engine
--fp16 启用FP16推理 --fp16
--int8 启用INT8推理 --int8
--iterations 推理迭代次数 --iterations=100
--avgRuns 用于计算平均延迟的轮次 --avgRuns=50
--dumpOutput 导出推理输出 --dumpOutput
总结一下:环境搭建这事,说白了就是版本匹配 + 路径配置。你只要把版本对应表记牢,安装顺序别搞反,环境变量配对,基本不会出大问题。trtexec这个工具,建议你多玩玩,它能帮你省掉很多写测试代码的时间。

嗯,这节课的内容就到这。下节课我们开始讲模型解析与网络定义,到时候会手把手带你写代码。记得先把环境搭好,不然下节课的代码你跑不起来。