环境搭建：Python、CUDA、PyTorch
说起环境配置，很多同学觉得这是小事。但我参与过几个蛋白质结构项目，吃过不少版本冲突的亏。环境搭不好，后面模型跑起来全是莫名其妙的错。今天就彻底把这部分理清楚。
AlphaFold 对硬件和库版本要求很具体。别看都叫 Python 3，差一个小版本，训练出来的结果可能完全不同。咱们一步步来。
版本选择与依赖关系
先列一张版本对应表。这是我反复验证过的组合，直接照抄就行。
| 组件 | 推荐版本 | 备注 |
|---|---|---|
| Python | 3.10.x | 别用 3.11+，有些依赖包还没兼容 |
| CUDA | 11.8 或 12.1 | 看显卡驱动；12.x 对 TensorRT 更友好 |
| cuDNN | 8.9.x（CUDA 11） 9.0.x（CUDA 12） |
和 CUDA 版本必须配对 |
| PyTorch | 2.1.x | 支持 torch.compile，导出 ONNX 方便 |
| JAX | 0.4.20 | AlphaFold2 依赖，CPU/GPU 版本分开装 |
为什么会这样选？Python 3.10 是主流框架支持最稳的版本。CUDA 12.1 配合 PyTorch 2.1 是目前导出 ONNX 的最佳拍档。我个人强烈建议优先考虑 12.1。
CUDA 与 cuDNN 安装
先检查你的显卡驱动。跑下面的命令：
nvidia-smi看右上角的 CUDA 版本。如果你是 545 以上的驱动，直接上 CUDA 12.1 没问题。低于 530 的驱动，老老实实装 CUDA 11.8。
安装 CUDA toolkit，我习惯用 runfile 方式。下载完执行：
chmod +x cuda_12.1.0_530.30.02_linux.run
sudo ./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_PATHcuDNN 需要先注册 NVIDIA 开发者账号。下载 tar 包后解压：
tar -xzvf cudnn-linux-x86_64-9.0.0.312_cuda12-archive.tar.xz
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*验证一下：
nvcc --version
cat /usr/local/cuda-12.1/include/cudnn_version.h | grep CUDNN_MAJOR -A 2两行都能输出具体版本号，说明装好了。
PyTorch 安装
这一步最折腾人。我建议直接用官方 Pip 命令，别去 conda 凑热闹。
pip install torch==2.1.2 torchvision==0.16.2 torchaudio==2.1.2 --index-url https://download.pytorch.org/whl/cu121注意 cu121 要和你的 CUDA 版本对应。如果装了 CUDA 11.8，把 cu121 换成 cu118。
装完验证 GPU 是否可用：
python -c "import torch; print(torch.cuda.is_available()); print(torch.__version__)"输出 True 和 2.1.2 就稳了。
JAX 安装（AlphaFold 专用）
AlphaFold2 的推理核心依赖 JAX。这里有个坑，JAX 的 GPU 版本需要手动指定。
pip install --upgrade "jax[cuda12_pip]==0.4.20" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html如果你是 AlphaFold2 的单体预测模式，这个版本足够了。多聚体预测要注意 jaxlib 的兼容性，我一般用 0.4.20 配 cuDNN 9.0。
其他依赖库
还有几个库是跑通 AlphaFold 和后续 ONNX 导出的必备品：
- haiku - DeepMind 的神经网络库，JAX 框架的封装
- dm-tree - 树形结构数据处理
- numpy - 科学计算基础
- scipy - 蛋白质结构计算用到的线性代数工具
安装命令：
pip install haiku dm-tree numpy scipy个人建议把所有依赖写进 requirements.txt 一次性装，省得遗漏。
整体流程梳理
这章内容不少，我画了一张图把整个搭建流程串起来，方便你对照。
验证环境完整性
所有组件装完后，跑一段综合测试脚本。确保各个库能正常调用 GPU。
import torch
import jax
import jaxlib
# PyTorch GPU 验证
print("PyTorch version:", torch.__version__)
print("CUDA available:", torch.cuda.is_available())
# JAX GPU 验证
print("JAX version:", jax.__version__)
print("JAX devices:", jax.devices())
print("JAX backend:", jax.lib.xla_bridge.get_backend().platform)
# 简单张量测试
x = torch.randn(3, 3).cuda()
y = torch.matmul(x, x)
print("PyTorch GPU tensor OK:", y.device)
z = jax.numpy.ones((3, 3))
print("JAX array OK")
如果输出里没有报错，CUDA 设备也正常显示，说明环境通了。
python -c "import jax; print(jax.devices())"环境搭建不是一次就完事的。后续在 ONNX 导出环节，可能还会遇到算子不兼容、TensorRT 版本冲突等问题。所以今天装的这些组件，最好在项目开始时就固定住版本号。说白了，就是给你的工程环境打个快照。
好了，这就是我个人积累下来的环境配置经验。跟着这份清单走，基本不会出错。配好环境之后，下一步就该进到真正的模型导出流程了。