用途
TorchDrug 是一个基于 PyTorch 的药物发现与分子科学机器学习工具箱。它将图神经网络、预训练模型和任务定义应用于分子、蛋白质和生物知识图谱，覆盖分子性质预测、蛋白质建模、知识图谱推理、分子生成、逆合成规划等方向，内置 40 多个精选数据集和 20 多种模型架构。
该技能适用于处理以下数据类型：SMILES 字符串或分子结构、蛋白质序列或 3D 结构（PDB 文件）、化学反应与逆合成、生物医学知识图谱、药物发现数据集。
主要能力
-
分子性质预测：从结构预测分子的化学、物理和生物性质，用于药物相似性与 ADMET 性质、毒性筛选、量子化学性质、结合亲和力预测。包含 20 多个分子数据集（BBBP、HIV、Tox21、QM9 等）、GNN 模型（GIN、GAT、SchNet）以及 PropertyPrediction 和 MultipleBinaryClassification 任务。
-
蛋白质建模：处理蛋白质序列、结构和性质，用于酶功能预测、蛋白质稳定性与溶解度、亚细胞定位、蛋白质-蛋白质相互作用、结构预测。包含 15 多个蛋白质数据集（EnzymeCommission、GeneOntology、PDBBind 等）、序列模型（ESM、ProteinBERT、ProteinLSTM）、结构模型（GearNet、SchNet）。
-
知识图谱推理：预测生物知识图谱中缺失的链接和关系，用于药物重定位、疾病机制发现、基因-疾病关联、多跳生物医学推理。包含通用知识图谱（FB15k、WN18）和生物医学知识图谱（Hetionet）、嵌入模型（TransE、RotatE、ComplEx）以及 KnowledgeGraphCompletion 任务。
-
分子生成：生成具有目标性质的新分子结构，用于从头药物设计、先导化合物优化、化学空间探索、性质引导生成。包含自回归生成、GCPN（基于策略的生成）、GraphAutoregressiveFlow 及性质优化工作流。
-
逆合成：从目标分子预测到起始原料的合成路线，用于合成规划、路线优化、合成可及性评估、多步规划。包含 USPTO-50k 反应数据集、CenterIdentification（反应中心预测）、SynthonCompletion（反应物预测）以及端到端 Retrosynthesis 流水线。
-
图神经网络模型：提供面向不同数据类型和任务的 GNN 架构目录，包括通用 GNN（GCN、GAT、GIN、RGCN、MPNN）、3D 感知模型（SchNet、GearNet）、蛋白质专用模型（ESM、ProteinBERT、GearNet）、知识图谱模型（TransE、RotatE、ComplEx、SimplE）、生成模型（GraphAutoregressiveFlow）。
-
数据集：40 多个精选数据集，涵盖化学、生物学和知识图谱，包括分子性质（药物发现、量子化学）、蛋白质性质（功能、结构、相互作用）、知识图谱（通用与生物医学）、逆合成反应。
使用方式
安装：
uv pip install torchdrug
或安装可选依赖
uv 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()
与 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 集成：
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)
适用场景
- 分子性质预测：预测候选药物的血脑屏障穿透性，加载 datasets.BBBP()，选择 GIN 模型，定义二分类 PropertyPrediction 任务，使用 scaffold 划分进行真实评估，以 AUROC 和 AUPRC 评估。
- 蛋白质功能预测：从序列预测酶功能，加载 datasets.EnzymeCommission()，选择 ESM（预训练）或 GearNet（带结构），定义多分类 PropertyPrediction 任务，微调预训练模型或从头训练。
- 基于知识图谱的药物重定位：在 Hetionet 中寻找新的疾病治疗方法，加载 datasets.Hetionet()，选择 RotatE 或 ComplEx，定义 KnowledgeGraphCompletion 任务，使用负采样训练，查询“Compound-treats-Disease”预测。
- 从头分子生成：生成针对目标结合优化的类药分子，训练性质预测器，选择 GCPN 进行基于强化学习的优化，定义结合亲和力、药物相似性、可合成性的奖励函数，生成候选并验证化学性质。
- 逆合成规划：为目标分子规划合成路线，加载 datasets.USPTO50k()，训练中心识别模型（RGCN）和合成子补全模型（GIN），组合为端到端逆合成流水线，递归应用于多步规划，检查构建块的商业可及性。
限制与常见问题
- 维度不匹配错误：检查 model.input_dim 是否与 dataset.node_feature_dim 匹配。
- 分子任务性能不佳：使用 scaffold 划分而非随机划分，尝试 GIN 而非 GCN。
- 蛋白质模型不学习：序列任务使用预训练 ESM，结构模型检查边构建。
- 大图内存错误：减小批量大小，使用梯度累积。
- 生成的分子无效：添加有效性约束，使用 RDKit 验证进行后处理。
官方文档：https://torchdrug.ai/docs/ GitHub：https://github.com/DeepGraphLearning/torchdrug 论文：TorchDrug: A Powerful and Flexible Machine Learning Platform for Drug Discovery