torchdrug

torchdrug

概述

TorchDrug 是一个基于 PyTorch 构建的综合性药物发现机器学习框架。它提供基于图的分子表示（原子作为节点，化学键作为边）、图神经网络（GNN）架构库、基准数据集和预训练模型，支持分子属性预测、药物-靶标相互作用、逆合成和分子生成设计等任务。TorchDrug 与 PyTorch Lightning 及标准机器学习工具集成，便于计算化学家和机器学习从业者使用。

适用场景

• 分子属性预测：训练或微调 GNN 模型，根据分子图预测 ADMET 属性（溶解度、毒性、渗透性）或生物活性（IC50、Ki）。
• 药物-靶标相互作用（DTI）预测：构建模型，预测化合物（SMILES）与蛋白质（序列或结构）之间的结合亲和力。
• 逆合成预测：使用基于模板或无模板的模型，为目标分子寻找合理的合成路线。
• 在大型分子数据集上预训练：利用在 ChEMBL 或 ZINC 上预训练的 GNN 表示，将其迁移学习到小型数据集。
• 分子生成：训练基于图的生成模型（GCPN、GraphAF），设计具有所需属性的新分子。
• 对 GNN 架构进行基准测试：比较 GraphConv、MPNN、GAT、AttentiveFP 在标准 MoleculeNet 任务上的表现。
• 如果要在不使用深度学习的情况下快速进行基于指纹的属性预测，请改用 RDKit + scikit-learn。
• 对于蛋白质结构任务（折叠、对接），请使用 ESMFold 或 DiffDock，而不是 TorchDrug。

前置条件

• Python 包：torchdrug、torch、torch-geometric、rdkit
• 环境：Python 3.8+，训练时建议使用兼容 CUDA 的 GPU
• 数据要求：SMILES 字符串或分子 SDF 文件；DTI 任务需要蛋白质序列

pip install torch torchvision --extra-index-url https://download.pytorch.org/whl/cu118
pip install torch-geometric
pip install torchdrug
pip install rdkit

快速开始

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

# Load a benchmark dataset and train a GNN for property prediction
dataset = datasets.BBBP("~/data/bbbp", node_feature="default", edge_feature="default")
print(f"Dataset: {len(dataset)} molecules, task: BBBP (blood-brain barrier penetration)")

# Define model: GIN encoder
model = models.GIN(
    input_dim=dataset.node_feature_dim,
    hidden_dims=[256, 256],
    short_cut=True,
    batch_norm=True,
    concat_hidden=True,
)

# Define training task
task = tasks.PropertyPrediction(
    model, task=dataset.tasks,
    criterion="bce", metric=("auprc", "auroc"),
)

# Train with the Solver
optimizer = torch.optim.Adam(task.parameters(), lr=1e-3)
solver    = core.Engine(task, dataset, None, None, optimizer, gpus=[0])
solver.train(num_epoch=50)
print("Training complete")

核心 API

模块 1：分子图表示

TorchDrug 将分子表示为带类型的图。data.Molecule 是核心数据结构。

from torchdrug import data
from rdkit import Chem

# Create a molecule from SMILES
smiles = "CC(=O)Oc1ccccc1C(=O)O"   # aspirin
mol    = data.Molecule.from_smiles(smiles, node_feature="default", edge_feature="default")

print(f"Atoms: {mol.num_node}")
print(f"Bonds: {mol.num_edge}")
print(f"Node feature dim: {mol.node_feature.shape}")  # (N_atoms, feature_dim)
print(f"Edge feature dim: {mol.edge_feature.shape}")  # (N_bonds*2, feature_dim)

# Convert a MoleculeNet / custom SMILES list to a dataset
from torchdrug import data as td_data
import pandas as pd

df = pd.read_csv("compounds.csv")   # columns: smiles, label
molecules = [td_data.Molecule.from_smiles(s) for s in df["smiles"] if s]
print(f"Loaded {len(molecules)} valid molecules")

# Check feature dimensions
print(f"Default atom feature dim: {molecules[0].node_feature.shape[1]}")

模块 2：GNN 架构

TorchDrug 提供 GIN、RGCN、GraphSAGE、GAT、MPNN、AttentiveFP 等模型。

from torchdrug import models, datasets

dataset = datasets.ESOL("~/data/esol", node_feature="default", edge_feature="default")
feature_dim = dataset.node_feature_dim

# Graph Isomorphism Network (GIN) — good default for property prediction
gin = models.GIN(
    input_dim=feature_dim,
    hidden_dims=[256, 256, 256],
    short_cut=True,
    batch_norm=True,
    concat_hidden=True,      # concatenate layer representations
)
print(f"GIN output_dim: {gin.output_dim}")

from torchdrug import models

# Message Passing Neural Network (MPNN) — captures edge features
mpnn = models.MPNN(
    input_dim=feature_dim,
    hidden_dim=256,
    edge_input_dim=16,       # edge feature dimension
    num_layer=4,
    num_gru_layer=1,
)

# Graph Attention Network (GAT) — attention-weighted neighbors
gat = models.GAT(
    input_dim=feature_dim,
    hidden_dims=[256, 256],
    edge_input_dim=16,
    num_head=8,
    batch_norm=True,
)
print(f"MPNN output_dim: {mpnn.output_dim}, GAT output_dim: {gat.output_dim}")

模块 3：分子属性预测

将 GNN 编码器与预测头组合，用于分类或回归。

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

# Regression example: ESOL aqueous solubility
dataset = datasets.ESOL("~/data/esol", node_feature="default", edge_feature="default")
train, val, test = dataset.split()
print(f"Train: {len(train)}, Val: {len(val)}, Test: {len(test)}")

model = models.GIN(
    input_dim=dataset.node_feature_dim,
    hidden_dims=[300, 300],
    short_cut=True,
    batch_norm=True,
    concat_hidden=True,
)

task = tasks.PropertyPrediction(
    model,
    task=dataset.tasks,       # list of property names
    criterion="mse",          # "mse" for regression, "bce" for classification
    metric=("mae", "rmse"),
    num_mlp_layer=2,
)

optimizer = torch.optim.Adam(task.parameters(), lr=1e-3, weight_decay=1e-5)
solver    = core.Engine(task, train, val, test, optimizer,
                        batch_size=32, log_interval=50)
solver.train(num_epoch=100)

# Evaluate on test set
metrics = solver.evaluate("test")
print(f"Test RMSE: {metrics['rmse']:.4f}")
print(f"Test MAE:  {metrics['mae']:.4f}")

模块 4：药物-靶标相互作用（DTI）预测

预测分子与蛋白质序列之间的结合亲和力。

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

# Load a DTI dataset (e.g., Davis kinase binding affinities)
dataset = datasets.Davis("~/data/davis",
                          mol_node_feature="default",
                          mol_edge_feature="default")
train, val, test = dataset.split()

# Molecule encoder
mol_model = models.GIN(
    input_dim=dataset.mol_node_feature_dim,
    hidden_dims=[256, 256],
    short_cut=True,
    batch_norm=True,
    concat_hidden=True,
)

# Protein encoder (CNN on sequence)
prot_model = models.ProteinCNN(
    input_dim=21,            # amino acid vocabulary size
    hidden_dims=[128, 128, 128],
    kernel_size=3,
)

task = tasks.InteractionPrediction(
    mol_model, prot_model,
    task=dataset.tasks,
    criterion="mse",
    metric=("rmse", "pearsonr"),
)

optimizer = torch.optim.Adam(task.parameters(), lr=1e-3)
solver    = core.Engine(task, train, val, test, optimizer,
                        batch_size=64, log_interval=100)
solver.train(num_epoch=50)

metrics = solver.evaluate("test")
print(f"DTI Test RMSE: {metrics['rmse']:.4f}")
print(f"DTI Pearson r: {metrics['pearsonr']:.4f}")

模块 5：逆合成预测

预测单步逆合成断键，寻找合理的构建单元。

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

# USPTO-50k retrosynthesis benchmark
dataset = datasets.USPTO50k("~/data/uspto50k",
                             as_synthon=False,
                             atom_feature="default",
                             bond_feature="default")
train, val, test = dataset.split()

# Reaction-predicting GNN
model = models.RGCN(
    input_dim=dataset.node_feature_dim,
    hidden_dims=[256, 256, 256],
    num_relation=dataset.num_bond_type,
    batch_norm=True,
)

task = tasks.CenterIdentification(
    model,
    feature=("graph", "atom", "bond"),
)

optimizer = torch.optim.Adam(task.parameters(), lr=1e-4)
solver    = core.Engine(task, train, val, test, optimizer,
                        batch_size=64, log_interval=100)
solver.train(num_epoch=50)
metrics = solver.evaluate("test")
print(f"Retrosynthesis top-1 accuracy: {metrics.get('accuracy', 'N/A')}")

模块 6：预训练模型和迁移学习

使用 TorchDrug 的预训练 GNN 表示作为下游任务的特征。

from torchdrug import models

# Load a GNN pretrained on ChEMBL with context-prediction self-supervised learning
pretrained_gin = models.GIN(
    input_dim=39,
    hidden_dims=[300, 300, 300, 300, 300],
    short_cut=False,
    batch_norm=True,
    concat_hidden=False,
)

# Load pretrained weights (download from TorchDrug model zoo)
import torch
ckpt = torch.load("gin_supervised_contextpred.pth", map_location="cpu")
pretrained_gin.load_state_dict(ckpt)
pretrained_gin.eval()

print(f"Pretrained GIN loaded, output_dim={pretrained_gin.output_dim}")
print("Use as encoder in PropertyPrediction task for transfer learning")

核心概念

基于图的分子表示

分子表示为带属性的图：原子是带有特征（原子序数、度、电荷、芳香性）的节点，化学键是带有特征（键类型、环成员关系）的边。所有 TorchDrug 模型都基于这些图表示运行，而不是基于 SMILES 字符串或指纹。

from torchdrug import data

mol = data.Molecule.from_smiles("c1ccccc1")   # benzene
print(f"Atoms: {mol.num_node}, Bonds: {mol.num_edge // 2}")
print(f"Atom features (first atom): {mol.node_feature[0]}")

Engine 和 Solver 模式

TorchDrug 使用 core.Engine（也称为 Solver）处理训练循环、日志记录、检查点保存和多 GPU 设置。将任务、训练/验证/测试数据集划分和优化器传递给 Engine，而不是手动编写训练循环。

# Engine handles: batch iteration, loss backward, logging, checkpointing
solver = core.Engine(
    task, train_set, valid_set, test_set, optimizer,
    batch_size=32,
    log_interval=100,
    gpus=[0, 1],         # multi-GPU support
)
solver.train(num_epoch=100)
solver.save("checkpoint.pth")

常见工作流

工作流 1：端到端 ADMET 属性预测

目标：训练 GIN 模型，根据 SMILES 预测血脑屏障穿透性，然后对新化合物进行预测。

import torch
import pandas as pd
from torchdrug import data, datasets, models, tasks, core

# 1. Load dataset
dataset = datasets.BBBP("~/data/bbbp", node_feature="default", edge_feature="default")
train, val, test = dataset.split()
print(f"BBBP: {len(train)} train, {len(val)} val, {len(test)} test molecules")

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

# 3. Train
optimizer = torch.optim.Adam(task.parameters(), lr=1e-3)
solver    = core.Engine(task, train, val, test, optimizer,
                        batch_size=32, log_interval=50)
solver.train(num_epoch=100)
metrics = solver.evaluate("test")
print(f"Test AUROC: {metrics['auroc']:.4f}")

# 4. Predict on new SMILES
new_smiles = ["CC(=O)Oc1ccccc1C(=O)O", "c1ccc(cc1)N"]
task.eval()
with torch.no_grad():
    for smi in new_smiles:
        mol = data.Molecule.from_smiles(smi, node_feature="default", edge_feature="default")
        batch = data.Batch.from_data_list([mol])
        pred  = task.predict(batch)
        print(f"  {smi}: BBB penetration probability = {pred.sigmoid().item():.3f}")

工作流 2：在 Tox21 上进行多任务属性预测

目标：使用共享的 GNN 编码器，同时预测 12 个毒性终点。

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

# Tox21: 12 toxicity assays, multi-label classification
dataset = datasets.Tox21("~/data/tox21", node_feature="default", edge_feature="default")
train, val, test = dataset.split()
print(f"Tox21 tasks ({len(dataset.tasks)}): {dataset.tasks}")

model = models.GIN(
    input_dim=dataset.node_feature_dim,
    hidden_dims=[300, 300, 300],
    short_cut=True, batch_norm=True, concat_hidden=True,
)

# Multi-task: one output head per toxicity assay
task = tasks.PropertyPrediction(
    model, task=dataset.tasks,
    criterion="bce",
    metric=("auroc",),
    num_mlp_layer=2,
)

optimizer = torch.optim.Adam(task.parameters(), lr=1e-3)
solver    = core.Engine(task, train, val, test, optimizer, batch_size=64)
solver.train(num_epoch=100)

metrics = solver.evaluate("test")
for name, val_score in metrics.items():
    print(f"  {name}: {val_score:.4f}")

关键参数

参数	模块	默认值	范围 / 选项	作用
hidden_dims	GIN/MPNN/GAT	[256, 256]	int 列表	GNN 层的宽度和深度
short_cut	GIN	False	True、False	在层之间添加残差连接
batch_norm	GIN/MPNN	False	True、False	在每一层之后应用批归一化
concat_hidden	GIN	False	True、False	将所有层的输出拼接为最终表示
num_mlp_layer	PropertyPrediction	1	1–4	GNN 后 MLP 预测头的深度
criterion	PropertyPrediction	"mse"	"mse"、"bce"、"ce"	损失函数：回归、二元 / 多标签分类
batch_size	Engine	32	8–512	训练批次大小

最佳实践

1. 在小型数据集上使用 concat_hidden=True 进行 GIN 训练：拼接所有层的输出能提供更丰富的分子表示，并且在训练数据有限（<10,000 个分子）时通常能提升性能。

2. 设置 batch_norm=True 以提升训练稳定性：批归一化可以降低模型对学习率和初始化的敏感度，尤其适用于较深的 GNN（3 层以上）。

3. 小型数据集从预训练 GNN 权重开始：TorchDrug 的模型库提供了通过自监督学习在 ChEMBL 上预训练的 GIN。在少于 1,000 个分子的数据集上，使用这些权重进行微调的效果优于随机初始化。

4. 在 scaffold 划分上验证，而非随机划分：随机训练 / 测试划分会高估泛化能力，因为结构相似的分子会同时出现在两个集合中。使用 dataset.split(test_scaffold_ratio=0.1) 可进行更真实的评估。

5. 处理多任务数据集中的缺失标签：许多 MoleculeNet 数据集（Tox21、SIDER）都存在缺失的测定值。TorchDrug 的 PropertyPrediction 任务会自动处理 NaN 标签，但请确认稀有测定任务的缺失率没有过高。

常见用法

用法：生成用于聚类的分子嵌入

适用场景：在嵌入空间中可视化分子库，或在 scikit-learn 模型中使用 GNN 特征。

import torch
import numpy as np
from torchdrug import data, models

model = models.GIN(input_dim=39, hidden_dims=[300, 300], concat_hidden=True)
model.eval()

smiles_list = ["CC(=O)O", "c1ccccc1", "CCN", "CC(=O)Oc1ccccc1C(=O)O"]
embeddings  = []
with torch.no_grad():
    for smi in smiles_list:
        mol   = data.Molecule.from_smiles(smi, node_feature="default")
        batch = data.Batch.from_data_list([mol])
        graph_feat = model(batch, batch.node_feature.float())["graph_feature"]
        embeddings.append(graph_feat.squeeze(0).numpy())

emb_matrix = np.stack(embeddings)
print(f"Embedding matrix: {emb_matrix.shape}")   # (N_mols, embed_dim)

用法：从 CSV 创建自定义分子数据集

适用场景：使用专有测定数据而非基准数据集进行训练。

from torchdrug import data
import torch

class CustomDataset(data.MoleculeDataset):
    def __init__(self, csv_path, smiles_col="smiles", label_col="activity"):
        import pandas as pd
        df = pd.read_csv(csv_path).dropna(subset=[smiles_col])
        smiles_list = df[smiles_col].tolist()
        targets     = df[label_col].tolist()
        self.load_smiles(smiles_list, {"activity": targets},
                         node_feature="default", edge_feature="default")
        self.tasks = ["activity"]

dataset = CustomDataset("assay_data.csv", smiles_col="smiles", label_col="pIC50")
print(f"Custom dataset: {len(dataset)} molecules")

故障排查

问题	原因	解决方案
ImportError: torchdrug	未安装软件包	安装 PyTorch 后运行 pip install torchdrug
CUDA error: device-side assert	标签数据类型不匹配	确保回归标签为 float，分类标签为 long
小型数据集上的测试指标较差	过拟合	使用预训练权重、添加 dropout 或降低模型深度
dataset.tasks 中出现 KeyError: task name	任务名称不匹配	打印 dataset.tasks 查看确切的任务名称；向 PropertyPrediction 传入相同的列表
RuntimeError: Expected all tensors on same device	CPU/GPU 张量混用	使用 solver = core.Engine(..., gpus=[0]) 确保设备放置一致
训练缓慢	仅使用 CPU 模式	安装兼容 CUDA 的 PyTorch；在 Engine 中设置 gpus=[0]
缺失的测定值导致 NaN 损失	数据集存在缺失标签	设置 criterion="bce" — TorchDrug 会在计算损失时屏蔽 NaN 标签

相关技能

• rdkit — 在使用 TorchDrug 前进行分子指纹计算和化学信息学预处理
• diffdock — 基于结构的对接，可补充 TorchDrug 的基于配体的预测

参考资料

• TorchDrug 文档 — 官方文档和教程
• TorchDrug GitHub (DeepGraphLearning/torchdrug) — 源代码
• Zhu 等人（2022），arXiv — TorchDrug 论文 — 平台原始论文
• MoleculeNet 基准 — TorchDrug 使用的标准化数据集