117

OpenFold:开源 AlphaFold 复现项目怎么用

OpenFold 是 AlphaFold2 的可训练开源复现。这篇讲清它的价值、微调能力与实际使用场景。

OpenFold(Ahdritz 等,Nature Methods 2024)是 AlphaFold2 的完整开源复现,最重要的是它提供了训练代码——DeepMind 只发布了推理代码与权重,OpenFold 让社区第一次能够从头训练与微调这类模型。

它解决了什么问题

官方 AlphaFold2 OpenFold
推理 可用 可用(结果基本一致)
训练代码 不提供 提供
微调 不可行 可行
框架 JAX PyTorch
显存优化 一般 更好(可处理更长序列)
许可 Apache 2.0(代码);权重受限 Apache 2.0,含自训练权重
可修改性 模块化,易改

论文中的重要发现

  • 训练动力学的分析:OpenFold 论文详细记录了训练过程中模型能力的涌现顺序——先学会二级结构,再学会长程接触,最后学会整体折叠。这个观察对理解模型如何工作很有价值;
  • 数据量的影响用远少于 PDB 全量的数据(如仅千分之几)训练,仍能达到接近的性能。这说明 AF2 的成功不完全依赖数据规模,架构设计的贡献很大;
  • MSA 的重要性验证:系统的消融实验量化了 MSA 深度对准确度的影响;
  • OpenProteinSet:论文同时发布了预计算的 MSA 数据集,省去了社区重复做最耗时的 MSA 搜索——这个数据集本身的价值不低于代码。

使用

# 安装(建议用官方 conda 环境)
git clone https://github.com/aqlaboratory/openfold.git
cd openfold
conda env create -f environment.yml
conda activate openfold
python setup.py install

# 下载权重
bash scripts/download_openfold_params.sh openfold/resources
# 也可以用官方 AF2 权重(推理结果与官方一致)
bash scripts/download_alphafold_params.sh openfold/resources

# 推理
python run_pretrained_openfold.py \
    fasta_dir/ \
    data/pdb_mmcif/mmcif_files/ \
    --uniref90_database_path data/uniref90/uniref90.fasta \
    --mgnify_database_path data/mgnify/mgy_clusters.fa \
    --pdb70_database_path data/pdb70/pdb70 \
    --uniclust30_database_path data/uniclust30/uniclust30_2018_08 \
    --output_dir output/ \
    --model_device cuda:0 \
    --config_preset model_1_ptm

# 用预计算的 MSA(跳过最慢的搜索步骤)
python run_pretrained_openfold.py \
    fasta_dir/ \
    data/pdb_mmcif/mmcif_files/ \
    --use_precomputed_alignments alignments/ \
    --output_dir output/

微调:OpenFold 的独特价值

# 场景:你有一批特定类型蛋白的实验结构
#      (如某个蛋白家族、抗体、膜蛋白)
#      想让模型在这类蛋白上表现更好
#
# 官方 AF2 做不到这一点 —— 没有训练代码

python train_openfold.py \
    mmcif_dir/ \
    alignment_dir/ \
    template_mmcif_dir/ \
    output_dir/ \
    2021-10-10 \
    --template_release_dates_cache_path cache.json \
    --precision bf16 \
    --gpus 8 \
    --replace_sampler_ddp=True \
    --seed 42 \
    --deepspeed_config_path deepspeed_config.json \
    --resume_from_ckpt openfold/resources/openfold_params/finetuning_2.pt

# 典型的微调场景:
#   1) 抗体结构预测
#      CDR 环是 AF2 的弱项(见 364)
#      → 用大量抗体结构微调,可能改善
#
#   2) 特定蛋白家族
#      如 GPCR、离子通道、激酶
#      → 让模型专注该家族的特点
#
#   3) 特定构象状态
#      如激酶的活性/非活性构象
#      → 用标注了状态的结构训练
#
#   4) 加入实验约束
#      如交联质谱、SAXS 数据
#      → 修改损失函数纳入这些约束
#
# 现实提醒:
#   微调需要相当的计算资源(多卡 GPU、数天到数周)
#   且需要足够的高质量结构数据
#   → 大多数团队用不上;但对有特定需求的团队价值很大

OpenProteinSet

  • 内容:PDB 中所有蛋白链的预计算 MSA,以及数百万条 UniClust 序列的 MSA;
  • 为什么重要MSA 搜索是整个流程中最耗时的部分(可占总时间的 90% 以上)。有了预计算的 MSA,训练与批量推理的成本大幅降低;
  • 获取:通过 AWS Open Data 公开,可直接下载或在云上挂载;
  • 用途:训练自己的结构预测模型、研究 MSA 的作用、加速批量推理。

什么时候用 OpenFold

场景 建议
偶尔预测几个结构 用 AlphaFold DB 或 ColabFold(见 118《ColabFold》
大批量推理 OpenFold(显存优化更好)
需要微调 OpenFold(唯一选择)
研究模型内部机制 OpenFold(PyTorch 易调试)
修改架构做新研究 OpenFold(模块化)
需要商用许可的权重 OpenFold 自训练权重(Apache 2.0)
预测复合物 Boltz-2 / Chai-1(见 121《Boltz-2 技术报告精读》122《Chai-1 技术报告精读》

相关的开源生态

  • OpenFold3:社区正在推进的 AF3 级复现,目标是提供可训练的全原子复合物预测;
  • Boltz-1/2(见 120《Boltz-1 技术报告精读》121《Boltz-2 技术报告精读》):MIT 许可的 AF3 级模型,代码与权重完全开放——对多数用户而言是比等待 OpenFold3 更实际的选择
  • Chai-1(见 122《Chai-1 技术报告精读》):另一个开放权重的 AF3 级模型;
  • 共同意义:这些项目让「AF3 级能力」不再被单一机构垄断,对产业界的可用性是决定性的

关键要点

  • 提供训练代码是 OpenFold 相对官方 AF2 的核心价值——微调成为可能;
  • 论文发现用远少于全量 PDB 的数据也能达到接近性能,说明架构贡献很大;
  • OpenProteinSet 的预计算 MSA 省去了最耗时的步骤,价值不低于代码本身;
  • 只是偶尔预测结构的话用 ColabFold;OpenFold 的价值在微调与研究。

延伸资源