2、环境搭建:安装CUDA、cuDNN、TensorRT、ONNX Runtime、PyTorch,验证安装是否成功。
好,咱们正式开始动手了。
环境搭建这事儿,说难不难,说简单吧,踩坑的人真不少。我见过太多同学卡在版本匹配上,一卡就是半天。所以这一章,我带你一步步走,把每个坑都提前填上。
2.1 版本匹配:这是最关键的
先别急着下载。你得搞清楚一件事:这些工具不是独立的,它们之间有严格的依赖关系。说白了,CUDA版本决定了cuDNN版本,cuDNN版本又限制了TensorRT版本,而PyTorch和ONNX Runtime也得跟着走。
我个人习惯的做法是:先定PyTorch版本,再倒推其他。比如我常用PyTorch 2.0,那CUDA就得选11.7或11.8。为什么?因为PyTorch官方编译的二进制包就是基于这些CUDA版本的。你装个CUDA 12.0,PyTorch 2.0反而用不了——嗯,这就是个典型的坑。
这里我整理了一个经过验证的版本组合,你直接照着来就行:
| 组件 | 推荐版本 | 备注 |
|---|---|---|
| CUDA | 11.8 | 稳定,兼容性好 |
| cuDNN | 8.9.7 | 对应CUDA 11.x |
| TensorRT | 8.6.1 | 支持ONNX解析 |
| PyTorch | 2.0.1 | CUDA 11.8版本 |
| ONNX Runtime | 1.15.1 | GPU版本 |
2.2 安装CUDA 11.8
去NVIDIA官网下载CUDA 11.8的runfile。我建议用runfile,别用deb包——deb包有时候会跟系统自带的驱动打架。
# 下载
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
安装时注意:不要勾选Driver。除非你确定系统里没有NVIDIA驱动,否则勾了容易出问题。我一般只选CUDA Toolkit和CUDA Samples。
装完后配置环境变量:
export PATH=/usr/local/cuda-11.8/bin:$PATH
export LD_LIBRARY_PATH=/usr/local/cuda-11.8/lib64:$LD_LIBRARY_PATH
验证一下:
nvcc --version
看到输出里有“release 11.8”就对了。
2.3 安装cuDNN 8.9.7
cuDNN需要注册NVIDIA账号才能下载。你下载那个“cuDNN Library for Linux (x86_64)”的tar包就行。
# 解压
tar -xzvf cudnn-linux-x86_64-8.9.7.29_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 /usr/local/cuda-11.8/lib64/libcudnn*
验证:
cat /usr/local/cuda-11.8/include/cudnn_version.h | grep CUDNN_MAJOR -A 2
看到版本号就对了。
2.4 安装TensorRT 8.6.1
TensorRT我推荐用tar包安装,别用pip——pip版本有时候会缺一些工具,比如trtexec。
# 下载TensorRT 8.6.1的tar包
tar -xzvf TensorRT-8.6.1.6.Linux.x86_64-gnu.cuda-11.8.tar.gz
# 移动到合适位置
sudo mv TensorRT-8.6.1.6 /usr/local/
# 配置环境变量
export LD_LIBRARY_PATH=/usr/local/TensorRT-8.6.1.6/lib:$LD_LIBRARY_PATH
export PATH=/usr/local/TensorRT-8.6.1.6/bin:$PATH
验证:
trtexec --version
看到版本信息就对了。
另外,记得安装Python的TensorRT包:
pip install /usr/local/TensorRT-8.6.1.6/python/tensorrt-8.6.1-cp39-none-linux_x86_64.whl
注意选对Python版本。我用的Python 3.9,所以选cp39的。
2.5 安装PyTorch 2.0.1
这个最简单,直接用pip:
pip install torch==2.0.1 torchvision==0.15.2 torchaudio==2.0.2 --index-url https://download.pytorch.org/whl/cu118
验证:
python -c "import torch; print(torch.__version__); print(torch.cuda.is_available())"
输出应该是“2.0.1”和“True”。
2.6 安装ONNX Runtime GPU版本
pip install onnxruntime-gpu==1.15.1
验证:
python -c "import onnxruntime; print(onnxruntime.__version__); print(onnxruntime.get_device())"
输出应该是“1.15.1”和“GPU”。
如果显示“CPU”,说明没装对。我建议你检查一下是不是装了onnxruntime(不带gpu后缀)——这两个包很容易搞混。
2.7 完整验证脚本
最后,我写了个小脚本,一键验证所有组件:
import torch
import onnxruntime
import tensorrt as trt
print("=== 环境验证 ===")
# PyTorch
print(f"PyTorch版本: {torch.__version__}")
print(f"CUDA可用: {torch.cuda.is_available()}")
if torch.cuda.is_available():
print(f"CUDA版本: {torch.version.cuda}")
print(f"GPU数量: {torch.cuda.device_count()}")
print(f"当前GPU: {torch.cuda.get_device_name(0)}")
# ONNX Runtime
print(f"ONNX Runtime版本: {onnxruntime.__version__}")
print(f"ONNX Runtime设备: {onnxruntime.get_device()}")
# TensorRT
print(f"TensorRT版本: {trt.__version__}")
print("=== 验证完成 ===")
跑一下,如果所有信息都正确显示,恭喜你,环境搭好了。
2.8 常见问题
我整理了几个高频问题,你遇到了可以先看看:
- nvcc找不到:检查PATH里有没有/usr/local/cuda-11.8/bin
- libcudart.so找不到:检查LD_LIBRARY_PATH里有没有/usr/local/cuda-11.8/lib64
- TensorRT导入报错:检查Python whl版本对不对,别装成CPU版本的
- ONNX Runtime显示CPU:确认装的是onnxruntime-gpu,不是onnxruntime
嗯,环境搭建就这些。你跟着走一遍,应该半小时就能搞定。下一章咱们开始实战——把PyTorch模型转成ONNX,然后跑通TensorRT推理。