第2章环境搭建

2、环境搭建：CUDA与cuDNN安装、TensorRT安装与验证、开发环境配置

说实话，环境搭建这一步，看着简单，但坑是真不少。我见过太多同学卡在这一步，明明代码写得挺好，结果跑不起来——一查，全是环境问题。所以这一章，咱们把地基打牢。

2.1 CUDA与cuDNN：GPU加速的基石

先说说CUDA。它是NVIDIA推出的并行计算平台。说白了，就是让GPU不光能打游戏，还能帮你做科学计算。TensorRT跑在GPU上，底层依赖的就是CUDA。

CUDA版本怎么选？

我个人习惯是：先看你的显卡驱动支持的最高CUDA版本，再看TensorRT要求什么版本，取交集。别一上来就装最新的，兼容性容易出问题。

⚠️ 避坑指南： 我曾经在项目里直接装了CUDA 12.x，结果发现TensorRT 8.x根本不支持。折腾了两天才发现是版本不匹配。所以，先查TensorRT官方文档，再决定CUDA版本。

安装步骤（以Ubuntu 20.04为例）：

下载CUDA 11.8 runfile（我推荐用runfile，比deb包更可控）
执行安装命令，注意不要安装驱动（如果你已经装好了）
配置环境变量到 ~/.bashrc

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

# 安装（注意取消勾选Driver）
sudo sh cuda_11.8.0_520.61.05_linux.run

# 配置环境变量
echo 'export PATH=/usr/local/cuda-11.8/bin:$PATH' >> ~/.bashrc
echo 'export LD_LIBRARY_PATH=/usr/local/cuda-11.8/lib64:$LD_LIBRARY_PATH' >> ~/.bashrc
source ~/.bashrc

# 验证
nvcc --version

验证成功的话，你会看到CUDA版本号。嗯，到这里CUDA就装好了。

cuDNN安装

cuDNN是CUDA的深度神经网络加速库。说白了，就是给深度学习算子做了手工优化。没有它，卷积运算会慢很多。

安装其实很简单：下载对应CUDA版本的cuDNN，然后解压复制到CUDA目录下。

# 下载cuDNN 8.9.0 for CUDA 11.x（需要NVIDIA开发者账号）
# 解压
tar -xzvf cudnn-linux-x86_64-8.9.0.131_cuda11-archive.tar.xz

# 复制文件到CUDA目录
sudo cp cudnn-*-archive/include/cudnn*.h /usr/local/cuda/include/
sudo cp -P cudnn-*-archive/lib/libcudnn* /usr/local/cuda/lib64/

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

# 验证
cat /usr/local/cuda/include/cudnn_version.h | grep CUDNN_MAJOR -A 2

💡 小技巧： 我建议用软链接管理cuDNN版本。这样以后升级cuDNN时，只需要改链接指向，不用动其他配置。

2.2 TensorRT安装与验证

TensorRT的安装方式有两种：deb包和tar包。我个人更推荐tar包，因为灵活，可以随时切换版本。

下载TensorRT

去NVIDIA官网下载TensorRT 8.6 GA版本。注意选择对应CUDA 11.x的版本。

# 下载TensorRT 8.6.1 for Ubuntu 20.04
wget https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/8.6.1/TensorRT-8.6.1.6.Ubuntu-20.04.x86_64-gnu.cuda-11.8.tar.gz

# 解压到指定目录
sudo tar -xzvf TensorRT-8.6.1.6.Ubuntu-20.04.x86_64-gnu.cuda-11.8.tar.gz -C /opt/

# 配置环境变量
echo 'export LD_LIBRARY_PATH=/opt/TensorRT-8.6.1.6/lib:$LD_LIBRARY_PATH' >> ~/.bashrc
echo 'export PATH=/opt/TensorRT-8.6.1.6/bin:$PATH' >> ~/.bashrc
source ~/.bashrc

安装Python包

TensorRT的Python接口是独立安装的。你想想看，如果你用Python做模型部署，这一步必不可少。

# 进入TensorRT的Python目录
cd /opt/TensorRT-8.6.1.6/python

# 根据你的Python版本选择对应的whl文件
pip install tensorrt-8.6.1-cp38-none-linux_x86_64.whl

# 验证Python接口
python3 -c "import tensorrt as trt; print(trt.__version__)"

如果输出了8.6.1，说明安装成功。

验证工具：trtexec

TensorRT自带一个叫trtexec的命令行工具。它可以快速测试模型转换和推理性能。我经常用它做基准测试。

# 测试一个简单的ONNX模型
trtexec --onnx=model.onnx --fp16 --workspace=1024

# 查看TensorRT版本信息
trtexec --version

✅ 验证清单：

nvcc --version → 显示CUDA 11.8
cat /usr/local/cuda/include/cudnn_version.h → 显示cuDNN 8.9
python3 -c "import tensorrt as trt; print(trt.__version__)" → 显示8.6.1
trtexec --version → 显示TensorRT版本

2.3 开发环境配置

环境搭好了，接下来配置开发环境。我建议用Docker，这样环境可以复用，换机器也不怕。

Docker镜像方案（推荐）

NVIDIA官方提供了TensorRT的Docker镜像。直接拉下来就能用，省去很多麻烦。

# 拉取TensorRT镜像
docker pull nvcr.io/nvidia/tensorrt:23.08-py3

# 运行容器，挂载代码目录
docker run --gpus all -it --rm \
  -v /path/to/your/code:/workspace \
  nvcr.io/nvidia/tensorrt:23.08-py3

# 进入容器后验证
python3 -c "import tensorrt as trt; print(trt.__version__)"

本地Python虚拟环境

如果你不想用Docker，也可以用conda或venv。我建议用conda，因为管理不同版本的CUDA工具链更方便。

# 创建conda环境
conda create -n tensorrt_env python=3.8
conda activate tensorrt_env

# 安装PyTorch（用于导出模型）
pip install torch torchvision --index-url https://download.pytorch.org/whl/cu118

# 安装ONNX
pip install onnx onnxruntime-gpu

# 安装TensorRT Python包（前面已经装过）
# 安装可视化工具
pip install netron

💡 我的习惯： 我会在项目根目录放一个requirements.txt，把所有依赖写进去。这样新同事加入时，直接pip install -r requirements.txt就完事了。

IDE配置

我用的是VS Code。配置好远程开发环境后，可以直接在容器里写代码、调试。

安装Remote - Containers插件
在项目根目录创建.devcontainer/devcontainer.json
指定TensorRT镜像作为开发环境

{
  "name": "TensorRT Dev",
  "image": "nvcr.io/nvidia/tensorrt:23.08-py3",
  "runArgs": ["--gpus", "all"],
  "mounts": [
    "source=/path/to/your/code,target=/workspace,type=bind"
  ],
  "customizations": {
    "vscode": {
      "extensions": [
        "ms-python.python",
        "ms-toolsai.jupyter"
      ]
    }
  }
}

配置好后，打开VS Code，点击左下角的绿色按钮，选择"Reopen in Container"。等几秒钟，你就进入了一个完整的TensorRT开发环境。

⚠️ 注意： 我曾经在Windows上用WSL2配置环境，结果TensorRT的Python包死活装不上。后来发现是WSL2的GPU穿透有问题。建议直接用Linux物理机或Docker。

2.4 环境验证：跑一个简单的TensorRT程序

环境配好了，咱们跑个程序验证一下。下面这个例子，创建一个简单的网络，用TensorRT推理。

import tensorrt as trt
import numpy as np

# 创建logger和builder
logger = trt.Logger(trt.Logger.WARNING)
builder = trt.Builder(logger)

# 创建网络定义
network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))

# 添加输入层
input_tensor = network.add_input(
    name="input",
    dtype=trt.float32,
    shape=(1, 3, 224, 224)
)

# 添加一个简单的卷积层
conv = network.add_convolution(
    input=input_tensor,
    num_output_maps=16,
    kernel_shape=(3, 3),
    kernel=np.random.randn(16, 3, 3, 3).astype(np.float32),
    bias=np.random.randn(16).astype(np.float32)
)
conv.stride = (1, 1)
conv.padding = (1, 1)

# 添加ReLU激活
relu = network.add_activation(
    input=conv.get_output(0),
    type=trt.ActivationType.RELU
)

# 标记输出
network.mark_output(relu.get_output(0))

# 构建引擎
config = builder.create_builder_config()
config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 20)  # 1MB
engine = builder.build_serialized_network(network, config)

# 创建运行时
runtime = trt.Runtime(logger)
engine = runtime.deserialize_cuda_engine(engine)

# 创建上下文
context = engine.create_execution_context()

# 准备输入数据
input_data = np.random.randn(1, 3, 224, 224).astype(np.float32)

# 分配GPU内存
output = np.empty((1, 16, 224, 224), dtype=np.float32)

# 执行推理
context.execute_v2(
    bindings=[
        int(input_data.ctypes.data),
        int(output.ctypes.data)
    ]
)

print("✅ TensorRT推理成功！输出形状:", output.shape)

如果程序正常运行，没有报错，说明你的环境完全OK。嗯，到这里，环境搭建就完成了。

下一章，咱们开始讲ONNX模型导出。这是TensorRT优化的第一步，也是很多同学容易踩坑的地方。到时候我会分享一些我在项目中总结的经验。

📚 课程目录

2、环境搭建：CUDA与cuDNN安装、TensorRT安装与验证、开发环境配置

2.1 CUDA与cuDNN：GPU加速的基石

2.2 TensorRT安装与验证

2.3 开发环境配置

2.4 环境验证：跑一个简单的TensorRT程序