002、环境搭建：Python、PyTorch与CUDA开发环境配置

昨天深夜，隔壁组的小王跑来找我，一脸愁容：“哥，我模型训到一半报错了，CUDA error: out of memory，可我显卡显存明明是够的啊。”我让他跑了个nvidia-smi，一看，好家伙，PyTorch 1.8配CUDA 11.1，驱动版本却还停在440。版本对不上，显存管理直接乱了套。这类环境问题，十个AI新手里有八个踩过坑。今天咱们就彻底把Python、PyTorch和CUDA这“三件套”的配置捋清楚。

环境配置的本质：版本对齐

很多人以为装个Anaconda、pip install torch就完事了，其实真正的难点在于版本对齐。CUDA驱动、CUDA Toolkit、PyTorch版本、Python版本，甚至操作系统版本，一环扣一环。没对齐，轻则性能损失，重则直接跑不起来。我的习惯是：先定PyTorch版本，再倒推其他依赖。

比如最近我在做量化部署，需要PyTorch 2.0以上的特性。查官方文档（一定要看！），发现PyTorch 2.0需要CUDA 11.7或11.8。那么整个链路就确定了：驱动版本至少满足CUDA 11.8要求（≥520.61.05），然后装CUDA Toolkit 11.8，最后用pip安装对应版本的PyTorch。

实操：从驱动到PyTorch

第一步，检查显卡驱动。

nvidia-smi

右上角会显示CUDA Version，那是驱动最高支持的CUDA版本，不代表你已经装了CUDA Toolkit。如果驱动太老，去官网下载对应显卡的Studio驱动（比Game驱动更稳定），覆盖安装即可。

第二步，安装CUDA Toolkit。
这里有个关键选择：用系统级安装，还是conda虚拟环境安装？我强烈推荐后者。系统里装多个CUDA Toolkit容易打架，conda环境隔离更干净。

conda create -n torch20 python=3.9
conda activate torch20
conda install cudatoolkit=11.8 -c nvidia

注意这里用的是cudatoolkit，不是从NVIDIA官网下载的那个完整CUDA。conda版本只包含必要的运行时库，够用，而且不会污染系统环境。

第三步，安装PyTorch。
别急着pip install torch！去官网复制命令。现在PyTorch官网的安装命令生成器做得不错，但要注意：如果你用conda安装了cudatoolkit，pip安装PyTorch时要选CUDA 11.8、pip安装方式。命令大概长这样：

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

这里cu118就是CUDA 11.8的缩写。别用conda install pytorch，conda源里的版本经常滞后，容易出幺蛾子。

第四步，验证安装。
写个测试脚本，别只用torch.cuda.is_available()就完事，那只能验证CUDA是否可见。真正的“全家桶”验证得这么写：

import torch

print(f"PyTorch版本: {torch.__version__}")
print(f"CUDA可用: {torch.cuda.is_available()}")
print(f"CUDA版本: {torch.version.cuda}")
print(f"当前设备: {torch.cuda.get_device_name(0)}")

# 跑个张量计算，真枪实弹试一下
x = torch.randn(3, 3).cuda()
print(f"张量计算正常: {x @ x.t()}")

# 再验证下cudnn，很多加速依赖它
print(f"cuDNN版本: {torch.backends.cudnn.version()}")

如果这里报错，大概率是版本没对齐。常见的一个坑是：显示CUDA可用，但跑张量计算就报错。那可能是CUDA Toolkit和PyTorch编译时的CUDA版本不匹配。

虚拟环境管理：别把所有东西都塞base里

见过不少人把所有包装在base环境，半年后项目一多，依赖冲突到没法收拾。我的习惯是：一个项目，一个环境。环境配置文件environment.yml必须存好，方便复现。分享一个我常用的模板：

name: torch20_project
channels:
  - pytorch
  - nvidia
  - conda-forge
  - defaults
dependencies:
  - python=3.9
  - cudatoolkit=11.8
  - pip
  - pip:
    - torch==2.0.1
    - torchvision==0.15.2
    - torchaudio==2.0.2
    - -f https://download.pytorch.org/whl/cu118

用conda env create -f environment.yml一键重建环境，干净利落。

那些年踩过的坑

1. “为什么我的GPU跑得比CPU还慢？”
   检查torch.backends.cudnn.benchmark = True开了没。这个选项允许cuDNN自动寻找最优卷积算法，对于固定输入尺寸的场景能提速不少。但输入尺寸变化频繁的话，反而会拖慢速度。

2. “显存明明够，却报OOM？”
   除了驱动版本问题，还要注意PyTorch的显存分配策略。有个隐藏的坑：torch.cuda.empty_cache()不能乱用，频繁调用会引发内存碎片。真正该做的是管理好张量生命周期，及时del变量，必要时用with torch.no_grad():减少中间缓存。

3. “装到一半报SSL错误？”
   那是pip源或者conda源的问题。国内用户记得换国内镜像源。但安装PyTorch和CUDA相关包时，建议用官方源，镜像源可能同步不及时。可以单独为这几个包指定官方源，其他包用镜像。

4. “import torch报DLL加载失败？”
   典型的环境变量打架。特别是Windows上，如果之前装过Visual Studio或者旧版CUDA，PATH里可能有冲突。这时候用conda虚拟环境的优势就体现出来了——环境隔离，变量干净。

个人经验：环境配置是门手艺

配环境不是一次性任务，而是持续的过程。我的建议是：

• 文档化每一步：每次配环境，顺手把命令和版本号记在项目的README里。将来复现或者迁移环境，能省下大把时间。
• 驱动别追最新：除非你需要最新CUDA版本的特性，否则驱动选一个经过验证的稳定版本。生产环境尤其如此，新驱动可能引入新bug。
• 留个回退方案：配新环境前，先把旧环境用conda env export > environment_backup.yml备份。万一新环境翻车，能快速退回工作状态。
• 理解底层依赖：别只满足于“跑起来”。花点时间搞清楚CUDA、cuDNN、PyTorch之间的依赖关系。下次出问题，你就能快速定位到是哪个环节掉了链子。

环境配得干净，实验跑得才稳。这步功夫下足了，后面调试模型时才能排除环境干扰，聚焦在真正的算法问题上。好了，现在可以去试试你的GPU能不能吐出第一个“Hello, CUDA”了。