torchdrug

TorchDrug

将 TorchDrug 用作模块化的 PyTorch 图学习技术栈：

1. 加载一个 datasets.* 数据集，
2. 选择一个 models.* 表征模型，
3. 将其封装在一个 tasks.* 目标中，
4. 使用 core.Engine 训练并评估。

当前官方文档和最新版本均为 0.2.1。对于较新的 Python 或 PyTorch 组合，应视为未经验证，而不是默认假定其兼容。

从版本检查开始

生成或调试代码之前，请检查环境：

python --version
python -c "import torch; print(torch.__version__)"
python -c "import torchdrug; print(torchdrug.__version__)"

TorchDrug 0.2.1 支持的版本矩阵如下：

• Python 3.7 至 3.10
• PyTorch 1.8 至 2.0
• Linux、Windows 或 macOS
• Apple Silicon：PyTorch 1.13 或更高版本，仅支持 CPU；不支持 MPS

如果项目使用 Python 3.11+ 或 PyTorch 2.1+，请创建兼容环境或明确测试源码构建。不要将这些组合描述为受支持。

安装

建议使用专用的 Python 3.10 环境，并固定 TorchDrug 版本：

uv venv --python 3.10
source .venv/bin/activate
uv pip install "torch==2.0.0"

按照官方安装页面的说明，安装与 PyTorch 和 CUDA 版本完全匹配的 torch-scatter 和 torch-cluster wheel。对于仅使用 CPU 的 PyTorch 2.0 环境，以下是一组可复现的 wheel 组合：

uv pip install "torch-scatter==2.1.1" "torch-cluster==1.6.1" \
  --find-links "https://data.pyg.org/whl/torch-2.0.0+cpu.html"
uv pip install "torchdrug==0.2.1"

不要在不同环境之间复制 CUDA wheel URL。请确保 PyTorch 版本、CUDA 构建、Python ABI 和平台都匹配。在 Apple Silicon 上，官方文档要求从源码构建 torch-scatter 和 torch-cluster；请固定经过审核的源码版本，并做好仅使用 CPU 执行的准备。

标准属性预测工作流

使用文档中的 ClinTox → GIN → PropertyPrediction → Engine 流程：

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

dataset = datasets.ClinTox("~/molecule-datasets/")
lengths = [int(0.8 * len(dataset)), int(0.1 * len(dataset))]
lengths.append(len(dataset) - sum(lengths))
train_set, valid_set, test_set = torch.utils.data.random_split(dataset, lengths)

model = models.GIN(
    input_dim=dataset.node_feature_dim,
    hidden_dims=[256, 256, 256, 256],
    short_cut=True,
    batch_norm=True,
    concat_hidden=True,
)
task = tasks.PropertyPrediction(
    model,
    task=dataset.tasks,
    criterion="bce",
    metric=("auprc", "auroc"),
)

optimizer = torch.optim.Adam(task.parameters(), lr=1e-3)
solver = core.Engine(
    task,
    train_set,
    valid_set,
    test_set,
    optimizer,
    batch_size=1024,
)
solver.train(num_epoch=100)
solver.evaluate("valid")

仅在有受支持的 CUDA 设备时添加 gpus=[0]。在 CPU 上执行时省略 gpus。

对于二分类，task.predict(batch) 返回 logits；需要概率时，请应用 torch.sigmoid。在 0.2.1 中，归一化后的回归预测会以原始目标值范围返回，这是相较于旧版本的一项重大变更。

选择官方工作流

分子属性预测

• 数据集：datasets.ClinTox、BBBP、Tox21、QM9 或其他有文档说明的分子数据集。
• 模型：从 models.GIN 开始；如果所选特征配置提供边特征，请使用 edge_input_dim。
• 任务：tasks.PropertyPrediction。
• 阅读分子属性预测。

分子自监督预训练

• InfoGraph：使用 models.InfoGraph(gin_model, separate_model=False)，并将其封装在 tasks.Unsupervised 中。
• 属性掩码：tasks.AttributeMasking(model, mask_rate=0.15)。
• 使用相同的编码器进行微调，然后在训练 tasks.PropertyPrediction 之前，通过 strict=False 加载检查点。
• 阅读分子属性预测。

分子生成

• 数据集：datasets.ZINC250k(..., kekulize=True, atom_feature="symbol")。
• GCPN：将 models.RGCN 编码器封装在 tasks.GCPNGeneration 中。
• GraphAF：将节点和边的 models.GraphAF 流模型封装在 tasks.AutoregressiveGeneration 中。
• 教程中支持的优化任务为 "qed" 和 "plogp"；优化标准为 "nll" 和／或 "ppo"。
• 阅读分子生成。

逆合成

• 创建两个同步的 datasets.USPTO50k 视图：使用反应模式进行反应中心识别，并使用 as_synthon=True 进行合成子补全。
• 分别训练 tasks.CenterIdentification 和 tasks.SynthonCompletion。
• 将训练好的任务与 tasks.Retrosynthesis 组合；不要将原始模型直接传递给端到端任务。
• 阅读逆合成。

知识图谱推理

• 嵌入工作流：datasets.FB15k237 → models.RotatE → tasks.KnowledgeGraphCompletion。
• 神经推理工作流：使用 fact_ratio=0.75 的 models.NeuralLP。
• 阅读知识图谱推理。

蛋白质建模

• 使用 data.Protein.from_sequence、from_pdb 或 from_molecule 构建蛋白质。
• 序列编码器包括 models.ESM、ProteinCNN、ProteinResNet、 ProteinLSTM 和 ProteinBERT；结构编码器包括 models.GearNet。
• 使用文档中介绍的图构建层，不要调用不存在的 protein.residue_graph() 便捷方法。
• 阅读蛋白质建模。

编写可靠 TorchDrug 代码的规则

1. 遵循 0.2.1 API。 官方文档并非持续更新的最新版文档。
2. 优先使用文档中介绍的特征名称。 使用 atom_feature、bond_feature、residue_feature 和 mol_feature；在相关数据集构造函数中，node_feature、edge_feature 和 graph_feature 是已弃用的别名。
3. 让 Engine 预处理任务。 如果组合预训练任务而不构造其求解器，请手动调用每个任务的 preprocess()。
4. 保持配对数据集划分同步。 对于逆合成，在划分反应数据集和合成子数据集之前，重置为相同的随机种子。
5. 使用 TorchDrug 整理批次。 使用 data.graph_collate 或 core.Engine；通用 PyTorch 整理函数不知道如何打包 TorchDrug 图。
6. 区分模型、任务和引擎的参数。 代码凭空杜撰的常见原因，是将任务选项传给模型，或在需要组合任务时传入原始模型。
7. 验证生成的化学结构。 将模型输出视为候选结构，而非经实验验证或可合成的化合物。

故障排除

安装或导入失败

将 Python、PyTorch、torch-scatter 和 torch-cluster 作为一个兼容组合进行检查。大多数故障源于二进制轮子不匹配、不受支持的 Python 版本，或尝试使用 MPS。

特征维度不匹配

根据已加载的数据集构建模型维度：

• dataset.node_feature_dim
• dataset.edge_feature_dim
• dataset.num_bond_type
• 知识图谱使用 dataset.num_entity 和 dataset.num_relation

不要硬编码从其他特征配置中复制的维度。

设备不匹配

对于受支持的 CUDA 执行，将 gpus=[0] 传给 core.Engine。手动预测时，先整理批次，再使用 utils.cuda 将整个嵌套批次移至目标设备。

检查点不匹配

重新创建相同的模型和特征配置。对于从预训练到微调的迁移，使用 strict=False 加载检查点中的 "model" 状态；对于完整的求解器，使用 solver.save() 和 solver.load()。

参考索引

• 核心概念和数据结构
• 数据集
• 模型和架构
• 分子性质预测和预训练
• 蛋白质建模
• 分子生成
• 逆合成
• 知识图谱推理

上游来源

• TorchDrug 0.2.1 文档
• 教程索引
• 安装
• 软件包参考
• TorchDrug 0.2.1 发布说明

Agent 操作流程

1. 检查环境。 确认工具包版本和输入格式（SMILES、SDF、PDB、CIF），并在计算前对分子进行规范化处理。
2. 明确输入。 根据数据或用户要求确认格式、标识符和参数。凡是会改变结果的值，宁可询问，也不要猜测。
3. 先运行小规模任务。 先处理少量分子或一个目标，并检查输出。
4. 使用上述说明和参考资料执行完整任务。
5. 验证结果。 检查无效结构、单位一致性以及预测模型的适用域。
6. 报告。 说明运行了什么（版本、命令、参数）、检查了什么，以及仍有哪些不确定之处。

如果发生这种情况	请这样做
分子无法解析或规范化	说明原因并跳过；绝不悄悄丢弃或修改结构。
这些说明中的函数、标志或端点在已安装版本中不存在	查阅已安装版本自身的文档（help()、--help、官方文档），调整做法并告知用户。绝不编造 API。
必需的输入、标识符或参数不明确	询问用户，或在运行前明确说明所采用的假设。

诚信规则

• 绝不捏造结果、参数、标识符、引文或统计数据。如果无法运行或验证某项内容，请如实说明。
• 计算预测是待验证的假设，不是实验结果；应明确标注。
• 此处的版本相关细节可能已过时：请对照已安装版本的官方文档进行确认。
• 对于会产生费用、占用共享 GPU 或云配额、涉及个人或患者数据，或无法撤销的操作，务必先征得同意。

相关技能

• deepchem：使用多样化的特征化器和预构建数据集进行分子机器学习。
• neurokit2：使用 NeuroKit2 构建或审核可复现的研究工作流，用于生理时间序列预处理、事件／间隔分析、多…
• rowan：Rowan 是一个云原生分子建模与药物化学工作流平台，提供 Python API。