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 12.0,然后死活装不上TensorRT 8.5。查了半天才发现,TensorRT 8.5根本不支持CUDA 12。最后全卸了重装,白白浪费两小时。

这里我整理了一个经过验证的版本组合,你直接照着来就行:

组件 推荐版本 备注
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

看到版本号就对了。

💡 小技巧:我习惯把cuDNN的lib路径加到LD_LIBRARY_PATH里,虽然理论上复制过去就行,但有时候程序会找不到。加一下更保险。

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”。

🔑 关键点:如果torch.cuda.is_available()返回False,别慌。先检查CUDA环境变量配没配对,再检查PyTorch版本是不是对应CUDA 11.8的。我遇到过有人装了CPU版本的PyTorch,那当然用不了GPU。

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("=== 验证完成 ===")

跑一下,如果所有信息都正确显示,恭喜你,环境搭好了。

⚠️ 注意:如果你用的是Windows,安装流程大同小异,但路径配置要小心。Windows的PATH环境变量里,路径别带空格,不然TensorRT会找不到dll。我有个学员就是装在“Program Files”里,折腾了一下午。

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推理。