TorchDrug 药物发现工具箱

要解决什么问题

药物发现与分子科学里，数据形态杂：SMILES、PDB、反应式、生物医学知识图谱。任务也分散——ADMET 与毒性筛选、蛋白质功能/结构、药物-靶点结合、de novo 设计、逆合成多步规划、图谱链接预测。若每个任务各自搭 PyTorch 数据管线、任务定义和评估协议，重复劳动多，且分子任务若用随机拆分容易高估泛化。

技能如何工作

TorchDrug 把图神经网络、预训练模型与统一任务接口放在同一套框架里，面向分子、蛋白质与知识图谱三类对象。

核心能力概览：

• 分子属性预测：PropertyPrediction、MultipleBinaryClassification；数据集如 BBBP、HIV、Tox21、QM9；模型如 GIN、GAT、SchNet
• 蛋白质建模：序列侧 ESM、ProteinBERT，结构侧 GearNet；数据集含 EnzymeCommission、GeneOntology、PDBBind
• 知识图谱推理：KnowledgeGraphCompletion；Hetionet、TransE / RotatE / ComplEx 等
• 分子生成：自回归、GCPN、GraphAutoregressiveFlow，可与属性预测器组成多目标奖励
• 逆合成：USPTO-50k；CenterIdentification + SynthonCompletion 拼成端到端流水线

典型步骤（以类药属性为例）： 加载 datasets.BBBP() → 选 GIN 等分子 GNN → 定义二分类任务 → 用 scaffold 拆分做评估 → 看 AUROC / AUPRC。复杂工作流在 skill 内按 references/*.md 分主题导航（数据集、模型、训练与陷阱）。

集成： 分子对象可与 RDKit 互转；蛋白质侧可接 AlphaFold / ESM 结构或表征；训练可包进 PyTorch Lightning。

适用边界与注意点

• 化学有效性、类药性过滤仍常依赖 RDKit 等后处理；生成任务需显式约束，否则无效结构增多。
• 大图或大批次易 OOM：需调 batch、梯度累积；model.input_dim 须与 dataset.node_feature_dim 一致。
• 分子 benchmark 优先骨架拆分；蛋白质序列任务通常受益于预训练 ESM，纯从零训练往往吃力。
• Skill 文档偏工作流与参考索引，不等同于替代 官方文档 与论文中的全部 API 细节。