DiffDock

模型介绍

DiffDock 是由 Corso 等人提出的蛋白-配体分子对接扩散模型，可实现但复合物分子对接、批量分子对接、评估等功能。

模型描述

DiffDock 模型将 docking 视为生成式建模问题，在给定蛋白受体和小分子配体后，通过 SO(3)/SE(3) 上的扩散过程同时建模配体的平移、旋转和扭转自由度，生成候选结合构象。

当前无权重及数据集，近期会上传魔搭，后续支持指令下载；

适用场景

使用说明

1. OneCode 使用

可通过 OneCode 在线环境体验智能化一键式 AI4S 编程：

点击体验智能化一键式 AI4S 编程

2. 手动安装使用

硬件要求

• 推荐使用 GPU 或 DCU 运行。
• CPU 可以用于连通性验证，但速度较慢。
• DCU 用户需要预先安装 DTK，建议使用 DTK 25.04.2 以上版本或与当前集群匹配的 OneScience 推荐版本。

软件要求

DiffDock 常用依赖包括 PyTorch、PyTorch Geometric、RDKit、OpenBabel、e3nn、torch-scatter、torch-cluster、NumPy、SciPy、tqdm、PyYAML 等。数据集评估中的 GNINA 能量最小化需要额外安装 gnina。

环境检测

• NVIDIA GPU：

nvidia-smi

• 海光 DCU：

hy-smi

快速开始

1. 安装运行环境

conda create -n onescience311 python=3.11 -y
conda activate onescience311
pip install onescience[bio] -i http://mirrors.onescience.ai:3141/pypi/simple/  --trusted-host mirrors.onescience.ai

#如果下述代码运行存在找不到库的情况，需要激活cuda，参考下列代码

source ${ROCM_PATH}/cuda/env.sh
export LD_LIBRARY_PATH="$CONDA_PREFIX/lib:$LD_LIBRARY_PATH"
export LD_LIBRARY_PATH="$CONDA_PREFIX/lib/python3.11/site-packages/fastpt/torch/lib:$LD_LIBRARY_PATH"

2. 下载模型包&数据库

# 如 ModelScope 页面上的模型 ID 大小写不同，请以实际上架名称为准
modelscope download --model OneScience/diffdock --local_dir ./diffdock
cd diffdock

训练权重及数据集

当前无权重及数据集，近期会上传魔搭，后续支持指令下载；

3. 检查包内文件

当前包不包含预训练权重，下面命令应只看到源码、配置和示例输入文件：

find . -maxdepth 3 -type f

若需要采样或评估，请先通过训练得到 checkpoint，或将外部 DiffDock score/confidence 权重放到本地目录，并保证权重目录中存在 model_parameters.yml。

4. Score 模型训练

训练前请将 configs/training.yml 中的数据路径改为本机路径，常用字段包括：

• data.pdbbind_dir：处理后的 PDBBind 数据目录。
• data.moad_dir：处理后的 MOAD 数据目录。
• data.split_train：训练集复合物列表。
• data.split_val：验证集复合物列表。
• runtime.log_dir：训练输出目录。

启动训练：

cd scripts
bash train.sh

训练成功后，输出目录通常包含：

• model_parameters.yml：模型结构与训练参数。
• best_model.pt：验证损失最优权重。
• best_ema_model.pt：EMA 权重。
• best_inference_epoch_model.pt 或 best_ema_inference_epoch_model.pt：开启 inference validation 时保存的权重。
• last_model.pt：最后一个 epoch 的训练状态。

5. 分子对接采样

采样需要一个已训练好的 score 模型目录，例如：

outputs/train/diffdock_cg_example/
├── model_parameters.yml
└── best_model.pt

请将 configs/sampling.yml 中的关键字段改为实际路径：

• model.model_dir：score 模型目录。
• model.ckpt：score checkpoint 文件名。
• confidence.confidence_model_dir：confidence 模型目录；不使用 rerank 时设为 null。
• input.protein_path：蛋白 PDB 路径，可使用 data/6o5u_protein_processed.pdb。
• input.ligand_description：SMILES 字符串或配体 SDF/MOL2 路径，可使用 data/6o5u_ligand.sdf。
• runtime.out_dir：采样输出目录。

启动单复合物采样：

cd scripts
bash infer.sh

6. CSV 批量采样

将 configs/sampling.yml 中 input.protein_ligand_csv 设置为 CSV 路径，并将单复合物字段按需设为 null。CSV 建议包含以下列：

complex_name,protein_path,ligand_description,protein_sequence

其中 ligand_description 可以是 SMILES，也可以是 SDF/MOL2 文件路径。

7. 数据集评估

评估前请将 configs/evaluate.yml 中的数据集、模型和输出目录改成实际路径：

python -m scripts.evaluate --config configs/evaluate.yml

若启用 gnina.gnina_minimize=true，请先确认当前环境中可以直接执行：

gnina --help

8. Confidence 模型训练

confidence 训练入口位于已安装的 OneScience 包中，本 ModelScope 包不重复复制该模块。完成 score 模型训练后，可在 OneScience 环境中运行：

python -m onescience.confidence.diffdock.confidence_train \
    --original_model_dir outputs/train/diffdock_cg_example \
    --data_dir /path/to/PDBBind_processed \
    --split_train /path/to/splits/timesplit_no_lig_overlap_train \
    --split_val /path/to/splits/timesplit_no_lig_overlap_val

数据准备

DiffDock 训练通常需要预处理后的 PDBBind 或 MOAD 数据。推荐将数据组织到统一根目录，例如：

${ONESCIENCE_DATASETS_DIR}/diffdock/
├── PDBBind_processed/
├── MOAD_processed/
└── splits/
    ├── timesplit_no_lig_overlap_train
    ├── timesplit_no_lig_overlap_val
    └── timesplit_test

划分文件为纯文本格式，每行一个复合物名称，需与数据目录中的子目录名对应。

OneScience 官方信息

引用与许可证

• DiffDock 原始代码使用 MIT License。本仓库保留来源说明，并面向 OneScience ModelScope 自动运行场景进行整理。
• 如果在科研工作中使用 DiffDock 结果，建议引用 DiffDock、DiffDock-L 原始论文和 OneScience 相关项目信息，并根据实际任务补充 PDBBind、MOAD、RDKit、OpenBabel、GNINA 等数据集或工具引用。

@inproceedings{corso2023diffdock,
  title={DiffDock: Diffusion Steps, Twists, and Turns for Molecular Docking},
  author={Corso, Gabriele and St{\"a}rk, Hannes and Jing, Bowen and Barzilay, Regina and Jaakkola, Tommi},
  booktitle={International Conference on Learning Representations},
  year={2023}
}