2. 开发环境搭建：Ubuntu + CUDA + cuDNN + PyTorch/TensorRT 环境配置、Docker 镜像制作

说实话，模型移植调试这件事，环境搭不好，后面全是坑。我见过太多团队在算法上花了几个月，结果一到部署现场，驱动版本对不上、CUDA 跑不起来、TensorRT 报错——整个项目直接卡住。所以这一章，咱们把地基打牢。

我个人习惯把环境搭建分成两条路：裸机安装和Docker 容器化。裸机适合开发调试，Docker 适合交付和复现。两条路我都会讲清楚，你根据项目阶段选就行。

2.1 硬件与操作系统选型

先别急着装软件，硬件得先确认。工业检测场景下，GPU 通常用 NVIDIA 的 Jetson 系列（边缘端）或者 RTX/Quadro 系列（服务器端）。我建议至少 8GB 显存，不然 YOLOv8 的 batch size 稍微大一点就 OOM。

操作系统我推荐 Ubuntu 20.04 LTS 或 22.04 LTS。20.04 更稳，22.04 对新硬件支持更好。嗯，这里要注意：如果你用 Jetson，必须用 NVIDIA 定制的 Ubuntu 镜像，别自己去装通用版。

2.2 驱动与 CUDA 安装

驱动是第一步。我踩过最大的坑就是先装了 CUDA 再装驱动，结果版本冲突，折腾了一整天。正确的顺序是：先装驱动 → 再装 CUDA → 最后 cuDNN。

驱动装好之后，验证一下：

nvidia-smi
# 你应该能看到 GPU 型号、驱动版本、CUDA 版本号

然后装 CUDA。我个人习惯用 CUDA 11.8，因为 PyTorch 和 TensorRT 对它的支持最成熟。下载 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

安装时记得 取消勾选 Driver，只装 CUDA Toolkit。装完后配置环境变量：

export PATH=/usr/local/cuda-11.8/bin:$PATH
export LD_LIBRARY_PATH=/usr/local/cuda-11.8/lib64:$LD_LIBRARY_PATH

验证 CUDA：

nvcc --version
# 输出应该显示 release 11.8

2.3 cuDNN 安装

cuDNN 是深度神经网络的加速库。说白了，没有它，卷积运算会慢好几倍。下载需要 NVIDIA 开发者账号，选跟 CUDA 11.8 匹配的版本。

安装其实很简单，解压后把文件拷贝到 CUDA 目录：

tar -xzvf cudnn-linux-x86_64-8.9.7.29_cuda11-archive.tar.xz
sudo cp cudnn-*-archive/include/cudnn*.h /usr/local/cuda-11.8/include/
sudo cp 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*

验证 cuDNN 版本：

cat /usr/local/cuda-11.8/include/cudnn_version.h | grep CUDNN_MAJOR -A 2

2.4 PyTorch 与 TensorRT 安装

PyTorch 安装我推荐用 pip，别自己从源码编译——除非你有特殊需求。CUDA 11.8 对应的 PyTorch 版本是 2.0.x 或 2.1.x：

pip install torch==2.1.0 torchvision==0.16.0 torchaudio==2.1.0 --index-url https://download.pytorch.org/whl/cu118

验证 PyTorch 能否调用 GPU：

python -c "import torch; print(torch.cuda.is_available()); print(torch.cuda.get_device_name(0))"
# 输出应为 True 和你的 GPU 型号

TensorRT 安装稍微麻烦一点。我建议用 tar 包安装，而不是 deb 包。deb 包会把库装到系统目录，容易跟其他版本冲突。tar 包可以自由切换版本。

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

然后安装 Python 绑定：

pip install /path/to/TensorRT-8.6.1.6/python/tensorrt-8.6.1-cp310-none-linux_x86_64.whl

验证 TensorRT：

python -c "import tensorrt as trt; print(trt.__version__)"
# 输出应为 8.6.1.6

2.5 Docker 镜像制作

裸机环境搭好了，但说实话，每次换机器或者给同事交付，都要重新配一遍，太痛苦了。所以我现在所有项目都走 Docker。

Docker 镜像的核心思路是：把环境固化成一个可复现的容器。我一般基于 NVIDIA 官方镜像 nvcr.io/nvidia/pytorch:23.10-py3 来构建，它已经包含了 CUDA、cuDNN 和 PyTorch。

这是我的 Dockerfile 模板：

FROM nvcr.io/nvidia/pytorch:23.10-py3

# 安装系统依赖
RUN apt-get update && apt-get install -y \
    libgl1-mesa-glx \
    libglib2.0-0 \
    libsm6 \
    libxext6 \
    libxrender-dev \
    libgomp1 \
    && rm -rf /var/lib/apt/lists/*

# 安装 TensorRT Python 绑定
COPY TensorRT-8.6.1.6 /opt/TensorRT-8.6.1.6
RUN pip install /opt/TensorRT-8.6.1.6/python/tensorrt-8.6.1-cp310-none-linux_x86_64.whl

# 设置工作目录
WORKDIR /workspace

# 复制项目代码
COPY . /workspace

CMD ["/bin/bash"]

构建镜像：

docker build -t defect-detection:v1.0 .

运行容器：

docker run --gpus all -it --rm \
    -v /path/to/data:/workspace/data \
    defect-detection:v1.0

2.6 环境验证全流程

环境搭完别急着跑模型，先跑一个完整的验证脚本。我每次都会写一个 check_env.py：

import torch
import tensorrt as trt

print(f"PyTorch version: {torch.__version__}")
print(f"CUDA available: {torch.cuda.is_available()}")
print(f"GPU count: {torch.cuda.device_count()}")
print(f"GPU name: {torch.cuda.get_device_name(0)}")
print(f"TensorRT version: {trt.__version__}")

# 简单张量运算测试
x = torch.randn(3, 3).cuda()
y = torch.randn(3, 3).cuda()
z = torch.matmul(x, y)
print(f"Tensor test passed: {z.shape}")

# TensorRT 构建器测试
builder = trt.Builder(trt.Logger(trt.Logger.WARNING))
network = builder.create_network()
print("TensorRT builder works fine")

如果这个脚本能跑通，恭喜你，环境搭好了。如果报错，别慌，先看错误信息——90% 的问题是 CUDA 版本不匹配或者路径没配好。

2.7 知识体系总览

下面这张图把整个环境搭建的流程串起来了。你跟着这个顺序走，基本不会出错。

这张图你保存下来，每次搭环境对着走一遍。我刚开始带团队的时候，就是靠这张图让新人三天内上手。说白了，环境搭建没什么技术难度，就是细心和耐心。