alterlab-torchdrug

TorchDrug

概述

TorchDrug 是一个基于 PyTorch 的综合机器学习工具箱，适用于药物发现和分子科学。它可将图神经网络、预训练模型和任务定义应用于分子、蛋白质和生物知识图谱，涵盖分子属性预测、蛋白质建模、知识图谱推理、分子生成和逆合成规划，并提供 40+ 个精选数据集和 20+ 种模型架构。

何时使用此技能

在处理以下内容时，应使用此技能：

数据类型：

• SMILES 字符串或分子结构
• 蛋白质序列或 3D 结构（PDB 文件）
• 化学反应和逆合成
• 生物医学知识图谱
• 药物发现数据集

任务：

• 预测分子属性（溶解度、毒性、活性）
• 蛋白质功能或结构预测
• 药物-靶点结合预测
• 生成新的分子结构
• 规划化学合成路线
• 在生物医学知识库中进行链接预测
• 在科学数据上训练图神经网络

库和集成：

• TorchDrug 是主要库
• 常与 RDKit 搭配用于化学信息学
• 兼容 PyTorch 和 PyTorch Lightning
• 可与 AlphaFold 和 ESM 集成以处理蛋白质

不适用的情况

场景	改用
如今构建通用自定义 GNN（使用活跃维护的库和当前版本的 Python/torch）	alterlab-torch-geometric
希望以更少代码使用现成的特征化工具、MoleculeNet 基准和预训练模型	alterlab-deepchem
只需将分子转换为自有模型使用的特征矩阵	alterlab-molfeat
只需获取带有官方划分的标注基准数据集（ADMET/DTI）	alterlab-pytdc
查询现有生物医学知识图谱，而非训练知识图谱模型	alterlab-primekg

开始使用

安装

# TorchDrug 0.2.1 (last release, Jul 2023) requires Python >=3.7,<3.11 and
# torch >=1.8. It will NOT solve on Python 3.11+ — pin an older interpreter:
uv venv --python 3.10
uv pip install torch                        # install torch first (see below)
# torch-scatter / torch-cluster ship only as sdists, so they compile against the
# installed torch — either install from the PyG wheel index or ensure a C++ toolchain:
uv pip install torch-scatter torch-cluster \
  --find-links "https://data.pyg.org/whl/torch-${TORCH_VERSION}+${CUDA}.html"  # e.g. cpu, cu121
uv pip install torchdrug==0.2.1

注意事项：

• 请先安装 torch；TorchDrug 和 torch-scatter/torch-cluster 会针对已安装的 torch 构建图操作，构建版本不匹配会导致导入失败。
• TorchDrug 自带 data.DataLoader、data.Graph 和 core.Engine——请使用这些，而不是直接使用 PyTorch 中对应的类（参见下方循环）。
• TorchDrug 依赖 rdkit-pypi，而该包在 2022.09 之后没有适用于 Python 3.11 以上版本的 wheel，这也是该技术栈固定使用较旧解释器的原因之一。
• 该项目已停止维护（自 2023 年 7 月发布 0.2.1 以来没有新版本）；对于新的 Python/torch 技术栈，优先选择 alterlab-torch-geometric 或 alterlab-deepchem。如果确实需要 TorchDrug 的任务/数据集抽象（例如 GearNet、其逆合成流程），则使用此技能。

快速示例

import torch
from torchdrug import data, datasets, models, tasks

# Load molecular dataset. MoleculeNet-style datasets such as BBBP have no .split()
# method (only datasets with predefined splits do) — split explicitly:
dataset = datasets.BBBP("~/molecule-datasets/")
lengths = [int(0.8 * len(dataset)), int(0.1 * len(dataset))]
lengths += [len(dataset) - sum(lengths)]
train_set, valid_set, test_set = data.ordered_scaffold_split(dataset, lengths)
# (or torch.utils.data.random_split(dataset, lengths) for a random split)

# Define GNN model
model = models.GIN(
    input_dim=dataset.node_feature_dim,
    hidden_dims=[256, 256, 256],
    edge_input_dim=dataset.edge_feature_dim,
    batch_norm=True,
    readout="mean"
)

# Create property prediction task
task = tasks.PropertyPrediction(
    model,
    task=dataset.tasks,
    criterion="bce",
    metric=["auroc", "auprc"]
)

# Train with a native PyTorch loop.
# preprocess() builds the task's MLP prediction head (and label statistics) — core.Engine
# calls it for you; in a manual loop call it BEFORE creating the optimizer, or the head
# is missing and its parameters are never optimized.
task.preprocess(train_set, valid_set, test_set)
# NOTE: use torchdrug.data.DataLoader (NOT torch.utils.data.DataLoader) — its
# default graph_collate packs data.Graph objects; the stock PyTorch collate cannot.
optimizer = torch.optim.Adam(task.parameters(), lr=1e-3)
train_loader = data.DataLoader(train_set, batch_size=32, shuffle=True)

for epoch in range(100):
    for batch in train_loader:
        # task.forward returns (loss, metric), not a bare tensor.
        loss, metric = task(batch)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

对于标准（非自定义）流程，core.Engine(task, train_set, valid_set, test_set, optimizer, batch_size=...) 会封装此循环并处理批处理和设备。需要完全控制训练步骤时，请使用上面的手动循环。

核心功能

1. 分子属性预测

根据分子结构预测其化学、物理和生物属性。

使用场景：

• 类药性和 ADMET 属性
• 毒性筛查
• 量子化学属性
• 结合亲和力预测

主要组件：

• 20+ 个分子数据集（BBBP、HIV、Tox21、QM9 等）
• GNN 模型（GIN、GAT、SchNet）
• PropertyPrediction 和 MultipleBinaryClassification 任务

参考： 请参阅 references/molecular_property_prediction.md，了解：

• 完整的数据集目录
• 模型选择指南
• 训练工作流和最佳实践
• 特征工程详情

2. 蛋白质建模

处理蛋白质序列、结构和性质。

用例：

• 酶功能预测
• 蛋白质稳定性和溶解度
• 亚细胞定位
• 蛋白质-蛋白质相互作用
• 结构预测

关键组件：

• 15+ 个蛋白质数据集（EnzymeCommission、GeneOntology、PDBBind 等）
• 序列模型（ESM、ProteinBERT、ProteinLSTM）
• 结构模型（GearNet、SchNet）
• 适用于不同预测层级的多种任务类型

参考： 查看 references/protein_modeling.md，了解：

• 蛋白质专用数据集
• 序列模型与结构模型
• 预训练策略
• 与 AlphaFold 和 ESM 的集成

3. 知识图谱推理

预测生物知识图谱中缺失的链接和关系。

用例：

• 药物再利用
• 疾病机制发现
• 基因-疾病关联
• 多跳生物医学推理

关键组件：

• 通用知识图谱（FB15k、WN18）和生物医学知识图谱（Hetionet）
• 嵌入模型（TransE、RotatE、ComplEx）
• KnowledgeGraphCompletion 任务

参考： 查看 references/knowledge_graphs.md，了解：

• 知识图谱数据集（包括包含 45k 个生物医学实体的 Hetionet）
• 嵌入模型对比
• 评估指标和协议
• 生物医学应用

4. 分子生成

生成具有理想性质的新型分子结构。

用例：

• 从头药物设计
• 先导化合物优化
• 化学空间探索
• 性质引导的生成

关键组件：

• 自回归生成
• GCPN（基于策略的生成）
• GraphAutoregressiveFlow
• 性质优化工作流

参考： 查看 references/molecular_generation.md，了解：

• 生成策略（无条件、有条件、基于骨架）
• 多目标优化
• 验证和筛选
• 与性质预测的集成

5. 逆合成

预测从目标分子到起始原料的合成路线。

用例：

• 合成规划
• 路线优化
• 合成可及性评估
• 多步规划

关键组件：

• USPTO-50k 反应数据集
• CenterIdentification（反应中心预测）
• SynthonCompletion（反应物预测）
• 端到端逆合成流程

参考： 查看 references/retrosynthesis.md，了解：

• 任务分解（中心识别 → 合成子补全）
• 多步合成规划
• 商业可购性检查
• 与其他逆合成工具的集成

6. 图神经网络模型

全面汇总适用于不同数据类型和任务的 GNN 架构。

可用模型：

• 通用 GNN：GCN、GAT、GIN、RGCN、MPNN
• 感知 3D 信息：SchNet、GearNet
• 蛋白质专用：ESM、ProteinBERT、GearNet
• 知识图谱：TransE、RotatE、ComplEx、SimplE
• 生成式：GraphAutoregressiveFlow

参考： 查看 references/models_architectures.md，了解：

• 详细的模型说明
• 按任务和数据集选择模型的指南
• 架构对比
• 实现技巧

7. 数据集

40+ 个精选数据集，涵盖化学、生物学和知识图谱。

类别：

• 分子性质（药物发现、量子化学）
• 蛋白质性质（功能、结构、相互作用）
• 知识图谱（通用和生物医学）
• 逆合成反应

参考： 查看 references/datasets.md，了解：

• 包含规模和任务信息的完整数据集目录
• 数据集选择指南
• 加载和预处理
• 划分策略（随机、骨架）

常见工作流

工作流 1：分子性质预测

场景： 预测候选药物的血脑屏障渗透性。

步骤：

1. 加载数据集：datasets.BBBP()
2. 选择模型：用于分子图的 GIN
3. 定义任务：使用二分类的 PropertyPrediction
4. 使用骨架划分进行训练，以获得真实的评估结果
5. 使用 AUROC 和 AUPRC 进行评估

导航： references/molecular_property_prediction.md → 数据集选择 → 模型选择 → 训练

工作流 2：蛋白质功能预测

场景： 根据序列预测酶的功能。

步骤：

1. 加载数据集：datasets.EnzymeCommission()
2. 选择模型：ESM（预训练）或 GearNet（结合结构）
3. 定义任务：使用多分类的 PropertyPrediction
4. 微调预训练模型或从头开始训练
5. 使用准确率和各类别指标进行评估

导航： references/protein_modeling.md → 模型选择（序列与结构）→ 预训练策略

工作流 3：通过知识图谱进行药物再利用

场景： 在 Hetionet 中寻找新的疾病疗法。

步骤：

1. 加载数据集：datasets.Hetionet()
2. 选择模型：RotatE 或 ComplEx
3. 定义任务：KnowledgeGraphCompletion
4. 使用负采样进行训练
5. 查询 “Compound-treats-Disease” 预测结果
6. 根据可信度和作用机制进行筛选

导航： references/knowledge_graphs.md → Hetionet 数据集 → 模型选择 → 生物医学应用

工作流 4：从头生成分子

场景： 生成针对靶点结合进行优化的类药分子。

步骤：

1. 使用活性数据训练性质预测器
2. 选择生成方法：使用 GCPN 进行基于 RL 的优化
3. 定义奖励函数，综合考虑亲和力、类药性和可合成性
4. 在性质约束下生成候选分子
5. 验证化学结构并按类药性筛选
6. 根据多目标评分排序

导航： references/molecular_generation.md → 条件生成 → 多目标优化

工作流 5：逆合成规划

场景： 为目标分子规划合成路线。

步骤：

1. 加载数据集：datasets.USPTO50k()
2. 训练中心识别模型（RGCN）
3. 训练合成子补全模型（GIN）
4. 将二者组合成端到端逆合成工作流
5. 递归应用以进行多步规划
6. 检查构建模块的商业可得性

导航： references/retrosynthesis.md → 任务类型 → 多步规划

集成模式

与 RDKit 集成

在 TorchDrug 分子和 RDKit 之间进行转换：

from torchdrug import data
from rdkit import Chem

# SMILES → TorchDrug molecule
smiles = "CCO"
mol = data.Molecule.from_smiles(smiles)

# TorchDrug → RDKit
rdkit_mol = mol.to_molecule()

# RDKit → TorchDrug
rdkit_mol = Chem.MolFromSmiles(smiles)
mol = data.Molecule.from_molecule(rdkit_mol)

与 AlphaFold/ESM 集成

使用预测的结构：

from torchdrug import data, layers
from torchdrug.layers import geometry

# Load AlphaFold predicted structure
protein = data.Protein.from_pdb("AF-P12345-F1-model_v4.pdb")

# Build a residue-level graph with sequential + spatial edges
graph_construction_model = layers.GraphConstruction(
    node_layers=[geometry.AlphaCarbonNode()],
    edge_layers=[
        geometry.SpatialEdge(radius=10.0, min_distance=5),
        geometry.SequentialEdge(max_distance=2),
    ],
    edge_feature="gearnet",
)
# GraphConstruction operates on a packed protein batch, not a bare Protein.
graph = graph_construction_model(data.Protein.pack([protein]))

与 PyTorch Lightning 集成

将任务封装起来，以便使用 Lightning 进行训练：

import torch
import pytorch_lightning as pl   # or: import lightning.pytorch as pl

class LightningTask(pl.LightningModule):
    def __init__(self, torchdrug_task):
        super().__init__()
        self.task = torchdrug_task      # call task.preprocess(...) before training

    def training_step(self, batch, batch_idx):
        loss, metric = self.task(batch)  # forward returns (loss, metric)
        return loss

    def validation_step(self, batch, batch_idx):
        pred = self.task.predict(batch)
        target = self.task.target(batch)
        return {"pred": pred, "target": target}

    def configure_optimizers(self):
        return torch.optim.Adam(self.parameters(), lr=1e-3)

技术详情

深入了解 TorchDrug 的架构：

核心概念： 请参阅 references/core_concepts.md，了解：

• 架构理念（模块化、可配置）
• 数据结构（Graph、Molecule、Protein、PackedGraph）
• 模型接口和 forward 函数签名
• 任务接口（predict、target、forward、evaluate）
• 训练工作流和最佳实践
• 损失函数和指标
• 常见问题与调试

速查表

选择数据集：

• 分子性质 → references/datasets.md → 分子部分
• 蛋白质任务 → references/datasets.md → 蛋白质部分
• 知识图谱 → references/datasets.md → 知识图谱部分

选择模型：

• 分子 → references/models_architectures.md → GNN 部分 → GIN/GAT/SchNet
• 蛋白质（序列）→ references/models_architectures.md → 蛋白质部分 → ESM
• 蛋白质（结构）→ references/models_architectures.md → 蛋白质部分 → GearNet
• 知识图谱 → references/models_architectures.md → KG 部分 → RotatE/ComplEx

常见任务：

• 性质预测 → references/molecular_property_prediction.md 或 references/protein_modeling.md
• 生成 → references/molecular_generation.md
• 逆合成 → references/retrosynthesis.md
• KG 推理 → references/knowledge_graphs.md

理解架构：

• 数据结构 → references/core_concepts.md → 数据结构
• 模型设计 → references/core_concepts.md → 模型接口
• 任务设计 → references/core_concepts.md → 任务接口

常见问题排查

问题：维度不匹配错误 → 检查 model.input_dim 是否与 dataset.node_feature_dim 匹配 → 参见 references/core_concepts.md → 必备属性

问题：分子任务表现不佳 → 使用骨架划分，而非随机划分 → 尝试使用 GIN，而非 GCN → 参见 references/molecular_property_prediction.md → 最佳实践

问题：蛋白质模型未能学习 → 序列任务使用预训练的 ESM → 检查结构模型的边构建 → 参见 references/protein_modeling.md → 训练工作流

问题：处理大型图时出现内存错误 → 减小批次大小 → 使用梯度累积 → 参见 references/core_concepts.md → 内存效率

问题：生成的分子无效 → 添加有效性约束 → 使用 RDKit 验证进行后处理 → 参见 references/molecular_generation.md → 验证与筛选

资源

官方文档： https://torchdrug.ai/docs/ GitHub： https://github.com/DeepGraphLearning/torchdrug 论文： TorchDrug：用于药物发现的强大而灵活的机器学习平台

总结

根据任务转到相应的参考文件：

1. 分子性质预测 → molecular_property_prediction.md
2. 蛋白质建模 → protein_modeling.md
3. 知识图谱 → knowledge_graphs.md
4. 分子生成 → molecular_generation.md
5. 逆合成 → retrosynthesis.md
6. 模型选择 → models_architectures.md
7. 数据集选择 → datasets.md
8. 技术细节 → core_concepts.md

每份参考文档都全面介绍相应领域，并提供示例、最佳实践和常见用例。