第二章：环境搭建——NVIDIA驱动、CUDA Toolkit、cuDNN与Docker GPU环境

说实话，环境搭建这事儿，看着简单，坑是真不少。我见过太多人花了一整天装驱动，结果PyTorch跑起来还是报错“CUDA not available”。嗯，咱们今天就把这事彻底捋清楚。

2.1 NVIDIA驱动安装：别小看这一步

驱动是GPU和操作系统之间的桥梁。你想想看，没有驱动，系统根本认不出显卡。我个人习惯用NVIDIA官方.run文件安装，而不是系统包管理器里的版本。为什么？因为官方驱动更新更快，对CUDA版本的兼容性更好。

安装步骤其实就三步：

1. 卸载旧驱动（如果有的话）
2. 禁用nouveau开源驱动
3. 运行.run文件安装

# 卸载旧驱动
sudo apt-get purge nvidia*
sudo apt-get autoremove

# 禁用nouveau
echo 'blacklist nouveau' | sudo tee /etc/modprobe.d/blacklist-nvidia-nouveau.conf
sudo update-initramfs -u

# 重启后安装
chmod +x NVIDIA-Linux-x86_64-*.run
sudo ./NVIDIA-Linux-x86_64-*.run

装完后用nvidia-smi验证。如果能看到GPU信息，驱动就搞定了。

2.2 CUDA Toolkit配置：版本匹配是关键

CUDA Toolkit是GPU计算的开发环境。说白了，它提供了编译器和库，让代码能跑在GPU上。这里有个关键点：CUDA Toolkit版本必须和驱动版本兼容。

我个人习惯用runfile方式安装CUDA Toolkit，而不是deb包。因为runfile可以自定义安装路径，方便多版本共存。我在项目中经常需要同时维护CUDA 11.8和12.1两个环境，runfile方式就特别灵活。

# 下载CUDA Toolkit runfile
wget https://developer.download.nvidia.com/compute/cuda/12.1.0/local_installers/cuda_12.1.0_530.30.02_linux.run

# 安装（注意不要安装驱动部分）
sudo sh cuda_12.1.0_530.30.02_linux.run --toolkit --silent --override

装完后配置环境变量：

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

2.3 cuDNN安装：深度学习加速器

cuDNN是NVIDIA针对深度神经网络优化的库。它提供了卷积、池化、归一化等操作的GPU加速实现。没有cuDNN，PyTorch和TensorFlow也能跑，但速度会慢很多。

安装其实很简单，解压后把文件复制到CUDA目录就行：

# 解压cuDNN
tar -xzvf cudnn-linux-x86_64-8.9.0.131_cuda12-archive.tar.xz

# 复制到CUDA目录
sudo cp cudnn-*-archive/include/cudnn*.h /usr/local/cuda-12.1/include/
sudo cp -P cudnn-*-archive/lib/libcudnn* /usr/local/cuda-12.1/lib64/
sudo chmod a+r /usr/local/cuda-12.1/include/cudnn*.h /usr/local/cuda-12.1/lib64/libcudnn*

2.4 Docker GPU环境：一劳永逸的方案

说实话，如果你经常要在不同机器上部署训练环境，Docker是最省心的方案。你想想看，每次换机器都要重新装驱动、配CUDA、装cuDNN，多麻烦。Docker GPU环境可以让你一次配置，到处运行。

首先安装nvidia-docker2：

# 添加NVIDIA Docker仓库
distribution=$(. /etc/os-release;echo $ID$VERSION_ID)
curl -s -L https://nvidia.github.io/nvidia-docker/gpgkey | sudo apt-key add -
curl -s -L https://nvidia.github.io/nvidia-docker/$distribution/nvidia-docker.list | sudo tee / /etc/apt/sources.list.d/nvidia-docker.list

sudo apt-get update
sudo apt-get install -y nvidia-docker2
sudo systemctl restart docker

然后拉取官方镜像：

# 拉取PyTorch官方镜像
docker pull nvcr.io/nvidia/pytorch:23.12-py3

# 运行容器
docker run --gpus all -it --rm nvcr.io/nvidia/pytorch:23.12-py3

2.5 PyTorch/TensorFlow GPU验证

环境搭好了，怎么确认GPU真的能用？我一般写个简单的脚本验证：

# PyTorch GPU验证
import torch
print(f"PyTorch版本: {torch.__version__}")
print(f"CUDA可用: {torch.cuda.is_available()}")
print(f"GPU数量: {torch.cuda.device_count()}")
print(f"当前GPU: {torch.cuda.get_device_name(0)}")

# 跑个小矩阵乘法测试
a = torch.randn(1000, 1000).cuda()
b = torch.randn(1000, 1000).cuda()
c = torch.mm(a, b)
print(f"矩阵乘法结果形状: {c.shape}")

# TensorFlow GPU验证
import tensorflow as tf
print(f"TensorFlow版本: {tf.__version__}")
print(f"GPU设备: {tf.config.list_physical_devices('GPU')}")

# 跑个小训练测试
with tf.device('/GPU:0'):
    a = tf.random.normal([1000, 1000])
    b = tf.random.normal([1000, 1000])
    c = tf.matmul(a, b)
    print(f"矩阵乘法结果形状: {c.shape}")

如果输出能看到GPU信息，恭喜你，环境搭好了。如果报错，别慌，按这个顺序排查：

• nvidia-smi能不能用？不能的话驱动有问题
• CUDA环境变量配对了没？echo $PATH看看
• cuDNN版本对不对？cat /usr/local/cuda/include/cudnn_version.h | grep CUDNN_MAJOR
• Docker容器里有没有加--gpus all？

2.6 知识体系总览

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

好了，环境搭好了，下一章咱们就可以开始真正搞GPU训练优化了。记住，环境搭建是基础，基础不牢，地动山摇。花点时间把这一步做扎实，后面会省很多事。