TorchDrug：面向药物发现的深度学习库

TorchDrug 是 MilaGraph（Jian Tang 组）开源的药物发现深度学习库。它的定位介于「通用图框架」和「现成流水线」之间：既提供分子/蛋白/知识图谱的数据结构与预置模型，又保留 PyTorch 的可改写性，适合做方法研究和快速原型。

安装

pip install torchdrug
# 需与已装的 torch / torch-scatter 版本匹配，建议先固定 torch 版本再装

依赖链对 PyTorch 与 torch-scatter 的版本较敏感，装不上时优先查这两者的兼容矩阵，这是最常见的失败点。

分层设计

一个最小例子

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

dataset = datasets.ClinTox("~/molecule-datasets/")
train, valid, test = dataset.split()

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, valid, test, optimizer, batch_size=128)
solver.train(num_epoch=100)
solver.evaluate("valid")

Molecule 与 Protein 是真正的差异点：它们把化学/生物语义（原子类型、键、残基、二级结构）内建进图对象，而不是让你自己拼特征矩阵。

覆盖面最广的地方

• 横跨模态：小分子、蛋白序列/结构、生物医学知识图谱在同一套 API 下，做跨模态任务（如 DTI、知识图谱辅助的靶点发现）时省掉大量胶水代码。
• 生成与逆合成内建：GCPN、GraphAF 等图生成模型和逆合成任务是现成的，这在同类库里不多见。
• 预训练范式齐全：属性掩码、边预测、上下文预测等自监督策略都有实现，适合复现分子预训练相关工作。

选型建议

• 选 TorchDrug：做方法研究、需要跨分子/蛋白/知识图谱、要复现图生成或知识图谱推理类工作。
• 选 PyG（见 176《PyTorch Geometric》）：要完全自定义消息传递、追新算子、进生产。PyG 生态更大、维护更活跃。
• 选 DeepChem（见 172《DeepChem》）：只做小分子性质预测，且想要最短路径跑通。
• 选 Chemprop（见 174《Chemprop》）：单纯要一个强的性质预测基线，不打算改架构。
• 注意维护节奏：TorchDrug 的更新频率不如 PyG，长期生产依赖前先评估社区活跃度与版本兼容性。

延伸资源

• 对照框架：172《DeepChem》、174《Chemprop》、176《PyTorch Geometric》、177《DGL-LifeSci》；
• 图神经网络基础：006《Stanford CS224W 图神经网络》；分子生成见「AI 模型」模块。