第4章 AlphaFold 源代码下载与配置
既然我们已经摸清了模型的结构和推理流程，接下来就该动手拉代码了。这个环节看似简单，但我见过不少团队因为少配了一个依赖、没选对分支，浪费了一整天。今天，我就把自己的踩坑经验摊开来聊聊。
咱们这一章的核心任务：把 AlphaFold 官方仓库克隆到本地、把环境搭好、让代码能跑通。别小看这三步，背后的依赖管理和版本兼容性问题，真能让人头大。
4.1 拉取源代码
首先得把官方仓库拉下来。我习惯用一个单独的目录来管理所有模型项目，比如 ~/projects/alphafold。你根据自己情况来就行。
cd ~/projects
git clone https://github.com/deepmind/alphafold.git
cd alphafold这里有个细节：克隆完成后，我建议你立刻查看当前所在的分支。
git branch -a默认 checkout 的是 main 分支。但我们要注意，现在官方已经迭代了好几版，有 v2.0、v2.1、v2.2 等。我个人推荐用 v2.1 或 v2.2，这两个版本的推理代码相对稳定，且对 ONNX 导出更友好。
选好版本后，切到对应 tag 或者分支：
git checkout v2.1.1为什么要强调这个？我曾经遇到过一个团队，他们直接从 main 分支拉代码，结果发现里面有些实验性的代码还没合并完，导致导入失败。所以，选一个稳定的 release tag，能省掉很多烦恼。
4.2 看一眼目录结构
代码下载完，先别急着安装。花几分钟熟悉一下仓库的布局，对后续调试有很大帮助。主要目录大概是这样的：
| 目录/文件 | 说明 |
|---|---|
| alphafold/ | 主包，包含模型、推理、数据管线代码 |
| alphafold/model/ | 模型定义、权重加载、配置 |
| alphafold/notebooks/ | Jupyter 示例，可以用来快速测试 |
| run_alphafold.py | 单序列/多序列批推理入口脚本 |
| environment.yml | Conda 环境定义，核心依赖锁 |
| Dockerfile | Docker 构建文件，服务器部署首选 |
我个人非常推荐先看看 alphafold/model/ 这个目录，里面能直观地看到各个模块（如 JAX 原生模型、config 定义）。这对接下来的导出操作很重要。
4.3 配置运行环境
这块是最容易出问题的。AlphaFold 官方的环境定义是基于 Conda 的，文件就是 environment.yml。
官方环境文件的大致内容：
name: alphafold
channels:
- conda-forge
- defaults
dependencies:
- python=3.8
- pip
- cudatoolkit=11.0
- pip:
- jax==0.2.12
- jaxlib==0.1.65
- tensorflow-cpu==2.5.0
- numpy
- scipy
- pandas
- hmmer
- hhsuite
- kalign你可能会问，为什么这里没有装 PyTorch？因为 AlphaFold 原生推理用的是 JAX + TensorFlow 的组合。但我们要做 ONNX 导出，后续还得额外装 PyTorch 和 onnxruntime。这里先按官方的来，先把环境跑通。
安装命令：
conda env create -f environment.yml
conda activate alphafold这个过程大概需要 5–15 分钟，取决于你的网络和机器性能。我在项目中遇到过一个问题：如果服务器上的 CUDA 驱动版本太高（比如 11.4+），和文件中指定的 cudatoolkit=11.0 可能会冲突。解决方法是把 cudatoolkit 的版本号注释掉，让 conda 自动选一个兼容的版本。
4.4 下载预训练参数和数据库
这是最耗时的步骤。AlphaFold 的权重文件大概有 2.3 GB，而 BFD 和 MGnify 等数据库加起来超过 400 GB。如果你是做 ONNX 优化，其实不需要全量数据库，只要模型参数就够了。
只下载模型参数：
mkdir -p params
cd params
wget https://storage.googleapis.com/alphafold/alphafold_params_v2.1.1.tar.bz2
tar -xvf alphafold_params_v2.1.1.tar.bz2解压后，你会看到如下文件：
params/
|-- params_model_1.npz
|-- params_model_2.npz
|-- params_model_3.npz
|-- params_model_4.npz
|-- params_model_5.npz每个文件对应一个模型的参数。如果你想减少内存占用，可以只保留前面几个模型。我在测试 ONNX 导出时，一般只用 model_1 和 model_2，这样能快速迭代。
4.5 验证环境是否搭建成功
最后一步，跑一个简单的测试，确保一切就绪。
测试命令：
python -c "import sys; sys.path.insert(0, '.'); from alphafold.model import model; print('Model import OK')"如果输出 Model import OK，说明核心依赖和模型导入都没问题。
你还可以再测试一下 JAX 能否调用 GPU：
python -c "import jax; print(jax.devices())"正常应该是会打印出类似 [GpuDevice(id=0), ...] 的信息。如果只显示 CPU，说明 CUDA 和 JAX 的绑定没配好。这时候，通常就是回去检查 jaxlib 和 cudatoolkit 的版本匹配情况。
我习惯把这些步骤整理成一个 checklist，每次换机器或者项目迁移时，照着走一遍就能快速验证。
| 检查项 | 验证方法 | 预期结果 |
|---|---|---|
| 代码版本 | git log --oneline -1 |
对应 release tag |
| Python 版本 | python --version |
3.8.x |
| JAX 可用 | python -c "import jax; jax.numpy.ones(1)" |
无报错 |
| GPU 识别 | python -c "import jax; print(len(jax.devices()))" |
≥ 1 |
| 参数文件 | ls params/ | wc -l |
至少 2 个 .npz |
配置环境这个事情，说白了就是一次配好、长期受益。我建议在项目一开始就把环境搞干净，不要依赖系统级全局装包，也不用 pip 强行装跑不动的版本。后面做 ONNX 导出时，JAX 和 PyTorch 的共存问题可能会再冒出来，但至少基础的 AlphaFold 推理环境是可靠的。
好了，环境配好、参数下载完毕，接下来你就可以自由操作——导入模型、跑推理测试，甚至开始研究如何把我们手上的 JAX 模型转成 ONNX 格式。这部分内容，我们一步一步来拆解。