torchdrug

TorchDrug - PyTorch 原生图神经网络药物发现工具

技能概述

TorchDrug 是一个基于 PyTorch 构建的机器学习工具箱，专为药物发现和分子科学研究设计，提供分子属性预测、蛋白质建模、知识图谱推理、分子生成和逆合成规划等 40+ 数据集和 20+ 模型架构的完整解决方案。

适用场景

1. 药物分子设计与筛选

通过图神经网络模型快速预测分子性质，包括溶解度、毒性、生物活性等 ADMET 属性，支持基于血脑屏障穿透性（BBBP）等真实数据集的药物候选筛选，大幅提升早期药物发现效率。

2. 蛋白质功能与结构预测

利用预训练模型（如 ESM）和结构感知模型（如 GearNet）进行酶功能预测、亚细胞定位、蛋白质-蛋白质相互作用分析，同时支持与 AlphaFold 预测结构的无缝集成。

3. 生物医学知识图谱推理

基于 Hetionet 等大规模生物医学知识图谱（包含 45,000+ 实体），使用 TransE、RotatE 等嵌入模型进行药物重定位、疾病机制发现和多跳推理，发现潜在的治疗新方案。

核心功能

1. 分子属性预测

集成 BBBP、HIV、Tox21、QM9 等 20+ 分子数据集，提供 GIN、GAT、SchNet 等多种图神经网络架构，支持二分类和多任务属性预测，通过骨架拆分策略确保模型评估的真实性。

2. 逆合成规划

基于 USPTO-50k 化学反应数据集，实现从目标分子到原料的合成路线预测，包含反应中心识别和合成子补全两阶段模型，支持多步合成路径规划和商业可用性检查。

3. 灵活的模型生态系统

提供 GCN、GAT、GIN、RGCN、MPNN 等通用 GNN，3D 感知模型 SchNet 和 GearNet，以及知识图谱嵌入模型 TransE 和 ComplEx，全部基于 PyTorch 原生实现，支持 PyTorch Lightning 训练加速。

常见问题

TorchDrug 适合什么类型的用户？

TorchDrug 主要面向药物发现研究人员、计算生物学家和化学信息学家。如果你需要构建自定义图神经网络架构进行药物发现、蛋白质建模或知识图谱推理，TorchDrug 是理想选择。如果你更关注预训练模型和多样化的特征提取器，DeepChem 可能更合适；如果重点是基准数据集，可以考虑 PyTDc。

如何开始使用 TorchDrug？

使用 uv pip install torchdrug 即可安装，也可使用 torchdrug[full] 安装完整依赖。首先从数据集加载开始（如 datasets.BBBP()），选择合适的模型（如 GIN 用于分子图），定义预测任务，然后使用标准 PyTorch 训练循环或 PyTorch Lightning 进行模型训练。

TorchDrug 与其他药物发现库如何配合？

TorchDrug 设计了良好的互操作性：可与 RDKit 进行分子格式转换（SMILES 与 TorchDrug 分子对象），支持 AlphaFold 和 ESM 的预测蛋白质结构，同时完全兼容 PyTorch 生态系统，可轻松集成到现有的深度学习工作流程中。