torchdrug

torchdrug

概要

TorchDrugは、PyTorchを基盤とする創薬向けの包括的な機械学習フレームワークです。グラフベースの分子表現（原子をノード、結合をエッジとして表現）、グラフニューラルネットワーク（GNN）アーキテクチャのライブラリ、ベンチマーク用データセット、分子特性予測、薬物標的相互作用、逆合成、分子生成設計などのタスク向けの事前学習済みモデルを提供します。TorchDrugはPyTorch Lightningや標準的なMLツールと統合されているため、計算化学者とML実務者の両方が利用できます。

使用する場面

• 分子特性予測：分子グラフからADMET特性（溶解性、毒性、透過性）や生物活性（IC50、Ki）を予測するGNNモデルの学習またはファインチューニング。
• 薬物標的相互作用（DTI）の予測：化合物（SMILES）とタンパク質（配列または構造）間の結合親和性を予測するモデルの構築。
• 逆合成予測：テンプレートベースまたはテンプレートフリーのモデルを使用して、標的分子の合成可能な経路を特定。
• 大規模分子データセットでの事前学習：ChEMBLまたはZINCで事前学習したGNN表現を活用し、小規模データセットに転移学習。
• 分子生成：グラフベースの生成モデル（GCPN、GraphAF）を学習し、望ましい特性を持つ新規分子を設計。
• GNNアーキテクチャのベンチマーク：標準的なMoleculeNetタスクでGraphConv、MPNN、GAT、AttentiveFPを比較。
• 深層学習を使わずに高速なフィンガープリントベースの特性予測を行う場合は、代わりにRDKit + scikit-learnを使用してください。
• タンパク質構造タスク（フォールディング、ドッキング）には、TorchDrugではなくESMFoldまたはDiffDockを使用してください。

前提条件

• 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：逆合成予測

実現可能なビルディングブロックを見つけるために、1段階の逆合成切断を予測します。

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. 小規模データセットの GIN では concat_hidden=True を使用します：すべてのレイヤーの出力を連結すると、より豊かな分子表現が得られ、学習データが限られている場合（分子数が 10,000 未満）に性能が向上することがよくあります。

2. 学習を安定させるために batch_norm=True を適用します：バッチ正規化により、特に深い GNN（3 レイヤー以上）で、学習率や初期化に対する感度が低下します。

3. 小規模データセットでは、事前学習済みの GNN 重みから始めます：TorchDrug のモデル動物園では、自己教師あり学習によって ChEMBL で事前学習された GIN を利用できます。これらの重みからファインチューニングすると、分子数が 1,000 未満のデータセットではランダム初期化よりも優れた性能が得られます。

4. ランダム分割ではなく、スキャフォールド分割で評価します：構造が類似した分子が両方のセットに含まれるため、ランダムな学習／テスト分割では汎化性能が過大評価されます。より現実的に評価するには、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	ラベルの dtype が一致しません	回帰ラベルが float、分類ラベルが long であることを確認します
小規模データセットでテスト指標が低い	過学習	事前学習済みの重みを使用する、ドロップアウトを追加する、またはモデルの深さを減らします
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 Documentation — 公式ドキュメントとチュートリアル
• TorchDrug GitHub (DeepGraphLearning/torchdrug) — ソースコード
• Zhu et al. (2022), arXiv — TorchDrug paper — プラットフォームの原著論文
• MoleculeNet Benchmark — TorchDrug で使用される標準化データセット