PyTorch 环境搭建全攻略：CUDA/cuDNN 配置与多版本管理技巧

一、环境搭建前的准备工作

1. 硬件兼容性检测

# 检查NVIDIA GPU型号
nvidia-smi

# 验证CUDA支持的Compute Capability
lspci | grep -i nvidia

# 查看CUDA版本兼容性矩阵
https://developer.nvidia.com/cuda-gpus

2. 系统环境要求

组件	推荐配置	最低要求
操作系统	Ubuntu 20.04 LTS	Windows 10/11
显卡驱动	NVIDIA 450.80.02+	NVIDIA 418.39+
Python	3.8-3.11	3.6+
CUDA Toolkit	11.7 (推荐)	10.2+
cuDNN	8.4.0 (与 CUDA 版本匹配)	7.6.5+

二、CUDA Toolkit 安装指南

1. Ubuntu 系统安装

# 添加CUDA仓库
wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/cuda-keyring_1.0-1_all.deb
sudo dpkg -i cuda-keyring_1.0-1_all.deb
sudo apt update

# 安装CUDA 11.7
sudo apt install cuda-11-7

# 设置环境变量
echo 'export PATH=/usr/local/cuda-11.7/bin${PATH:+:${PATH}}' >> ~/.bashrc
echo 'export LD_LIBRARY_PATH=/usr/local/cuda-11.7/lib64${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}' >> ~/.bashrc
source ~/.bashrc

2. Windows 系统安装

1. 从CUDA Toolkit Archive下载对应版本
2. 运行安装程序，选择自定义安装
3. 安装完成后添加环境变量：
   • C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v11.7\bin
   • C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v11.7\libnvvp

3. 版本验证

# 检查CUDA版本
nvcc --version

# 验证CUDA样例程序
cd /usr/local/cuda-11.7/samples/1_Utilities/deviceQuery
make
./deviceQuery

三、cuDNN 加速库配置

1. 下载与安装

# 注册NVIDIA开发者账号后下载：
# https://developer.nvidia.com/rdp/cudnn-download

# Ubuntu系统安装
tar -xzvf cudnn-linux-x86_64-8.4.0.27_cuda11.6.tgz
sudo cp cudnn-linux-x86_64-8.4.0.27_cuda11.6/include/cudnn*.h /usr/local/cuda/include
sudo cp -P cudnn-linux-x86_64-8.4.0.27_cuda11.6/lib/libcudnn* /usr/local/cuda/lib64
sudo chmod a+r /usr/local/cuda/include/cudnn*.h /usr/local/cuda/lib64/libcudnn*

# Windows系统安装
# 解压后复制文件到CUDA安装目录对应的include和lib目录

2. 验证测试

import torch

# 检查CUDA可用性
print(torch.cuda.is_available())  # 输出 True

# 验证cuDNN版本
print(torch.backends.cudnn.version())  # 输出 8400

四、PyTorch 多版本管理技巧

1. 使用 conda 管理环境

# 创建Python 3.9环境
conda create -n pytorch-env python=3.9
conda activate pytorch-env

# 安装PyTorch 1.13.1 + CUDA 11.7
conda install pytorch==1.13.1 torchvision==0.14.1 torchaudio==0.13.1 cudatoolkit=11.7 -c pytorch

# 安装夜间版PyTorch
conda install pytorch torchvision torchaudio pytorch-cuda=11.7 -c pytorch-nightly -c nvidia

2. 使用 pip 安装特定版本

# 安装稳定版
pip3 install torch==1.12.1+cu116 torchvision==0.13.1+cu116 torchaudio==0.12.1 --extra-index-url https://download.pytorch.org/whl/cu116

# 安装CPU版本
pip3 install torch==1.13.1+cpu torchvision==0.14.1+cpu torchaudio==0.13.1 --extra-index-url https://download.pytorch.org/whl/cpu

3. 版本切换实践

# 查看已安装版本
conda list | grep pytorch

# 切换环境
conda activate pytorch-env

# 卸载当前版本
conda remove pytorch torchvision torchaudio

五、常见问题解决方案

1. CUDA 版本不匹配

# 错误提示：
# RuntimeError: CUDA error: CUBLAS_STATUS_ALLOC_FAILED when calling `cublasCreate(handle)`

# 解决方案：
# 1. 检查CUDA版本与PyTorch兼容性
# 2. 清理旧版本CUDA残留文件
# 3. 重新安装匹配版本组合

2. cuDNN 未正确加载

# 错误提示：
# RuntimeError: cuDNN error: CUDNN_STATUS_EXECUTION_FAILED

# 解决方案：
# 1. 确认cuDNN文件权限
# 2. 检查CUDA和cuDNN版本匹配
# 3. 重新编译cuDNN库

3. 虚拟环境冲突

# 解决方案：
# 使用conda创建独立环境：
conda create -n pytorch-1.13 python=3.9
conda activate pytorch-1.13

六、高级环境管理技巧

1. Docker 容器化部署

# Dockerfile
FROM nvidia/cuda:11.7.1-cudnn8-devel-ubuntu20.04

WORKDIR /app

RUN apt-get update && apt-get install -y python3-pip
RUN pip3 install torch==1.13.1+cu117 torchvision==0.14.1+cu117 torchaudio==0.13.1

CMD ["python3", "main.py"]

2. 多版本共存方案

# 管理多个CUDA版本
export CUDA_HOME=/usr/local/cuda-11.7
export PATH=$CUDA_HOME/bin:$PATH
export LD_LIBRARY_PATH=$CUDA_HOME/lib64:$LD_LIBRARY_PATH

# 切换CUDA版本
ln -sf /usr/local/cuda-11.7 /usr/local/cuda

3. 性能优化配置

# 设置CUDA缓存大小
export CUDA_CACHE_MAXSIZE=10737418240

# 启用Tensor Cores加速
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True

七、环境验证脚本

import torch
import sys

def print_env_info():
    print(f"Python Version: {sys.version.split()[0]}")
    print(f"PyTorch Version: {torch.__version__}")
    print(f"CUDA Version: {torch.version.cuda}")
    print(f"cuDNN Version: {torch.backends.cudnn.version()}")
    print(f"Available GPUs: {torch.cuda.device_count()}")
    if torch.cuda.is_available():
        print(f"Current GPU: {torch.cuda.get_device_name(0)}")

if __name__ == "__main__":
    print_env_info()