TorchDrug药物发现工具箱Skill torchdrug

---

name: torchdrug description: “基于图的药物发现工具包。分子属性预测（ADMET）、蛋白质建模、知识图谱推理、分子生成、逆合成、GNNs（GIN、GAT、SchNet），40+ 数据集，用于基于 PyTorch 的分子、蛋白质和生物医学图的机器学习。”

TorchDrug

概述

TorchDrug 是一个全面的基于 PyTorch 的机器学习工具箱，用于药物发现和分子科学。应用图神经网络、预训练模型和任务定义到分子、蛋白质和生物医学知识图谱，包括分子属性预测、蛋白质建模、知识图谱推理、分子生成、逆合成规划，包含 40+ 精选数据集和 20+ 模型架构。

何时使用此技能

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

数据类型：

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

任务：

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

库和集成：

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

入门指南

安装

pip install torchdrug
# 或带有可选依赖
pip install torchdrug[full]

快速示例

from torchdrug import datasets, models, tasks
from torch.utils.data import DataLoader

# 加载分子数据集
dataset = datasets.BBBP("~/molecule-datasets/")
train_set, valid_set, test_set = dataset.split()

# 定义 GNN 模型
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"
)

# 创建属性预测任务
task = tasks.PropertyPrediction(
    model,
    task=dataset.tasks,
    criterion="bce",
    metric=["auroc", "auprc"]
)

# 使用 PyTorch 训练
optimizer = torch.optim.Adam(task.parameters(), lr=1e-3)
train_loader = DataLoader(train_set, batch_size=32, shuffle=True)

for epoch in range(100):
    for batch in train_loader:
        loss = task(batch)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

核心能力

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 获取：

• 任务分解（中心 ID → 合成子完成）
• 多步合成规划
• 商业可用性检查
• 与其他逆合成工具的集成

6. 图神经网络模型

用于不同数据类型和任务的 GNN 架构综合目录。

可用模型：

• 通用 GNNs：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. 选择生成方法：用于基于 RL 优化的 GCPN
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 分子
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

# 加载 AlphaFold 预测结构
protein = data.Protein.from_pdb("AF-P12345-F1-model_v4.pdb")

# 构建带空间边的图
graph = protein.residue_graph(
    node_position="ca",
    edge_types=["sequential", "radius"],
    radius_cutoff=10.0
)

与 PyTorch Lightning 集成

包装任务用于 Lightning 训练：

import pytorch_lightning as pl

class LightningTask(pl.LightningModule):
    def __init__(self, torchdrug_task):
        super().__init__()
        self.task = torchdrug_task

    def training_step(self, batch, batch_idx):
        return self.task(batch)

    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）
• 模型接口和前向函数签名
• 任务接口（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

每个参考提供其领域的全面覆盖，包含示例、最佳实践和常见用例。