alterlab-torchdrug

TorchDrug

概要

TorchDrug は、創薬と分子科学のための包括的な PyTorch ベースの機械学習ツールボックスです。グラフニューラルネットワーク、事前学習済みモデル、タスク定義を分子、タンパク質、生物医学ナレッジグラフに適用できます。対象タスクには、分子特性予測、タンパク質モデリング、ナレッジグラフ推論、分子生成、逆合成計画などがあり、厳選された 40 以上のデータセットと 20 以上のモデルアーキテクチャを備えています。

このスキルを使用する場面

次のような作業でこのスキルを使用してください。

データの種類：

• SMILES 文字列または分子構造
• タンパク質配列または 3D 構造（PDB ファイル）
• 化学反応と逆合成
• 生物医学ナレッジグラフ
• 創薬データセット

タスク：

• 分子特性（溶解度、毒性、活性）の予測
• タンパク質の機能または構造の予測
• 薬物と標的の結合予測
• 新しい分子構造の生成
• 化学合成経路の計画
• 生物医学ナレッジベースにおけるリンク予測
• 科学データを用いたグラフニューラルネットワークのトレーニング

ライブラリと統合：

• 主に TorchDrug を使用します
• ケモインフォマティクスには RDKit と併用されることがよくあります
• PyTorch および PyTorch Lightning と互換性があります
• タンパク質向けに AlphaFold および ESM と統合できます

該当しない場合

ケース	代わりに使用
現在の Python/torch で汎用のカスタム GNN を構築する（活発に開発されているライブラリ）	alterlab-torch-geometric
少ないコードで既製の特徴量化器、MoleculeNet ベンチマーク、事前学習済みモデルを使用する	alterlab-deepchem
独自モデル用の特徴量行列に分子を変換するだけの場合	alterlab-molfeat
公式の分割データを備えたラベル付きベンチマークデータセット（ADMET/DTI）を入手するだけの場合	alterlab-pytdc
KG モデルをトレーニングするのではなく、既存の生物医学ナレッジグラフを照会する場合	alterlab-primekg

はじめに

インストール

# TorchDrug 0.2.1 (last release, Jul 2023) requires Python >=3.7,<3.11 and
# torch >=1.8. It will NOT solve on Python 3.11+ — pin an older interpreter:
uv venv --python 3.10
uv pip install torch                        # install torch first (see below)
# torch-scatter / torch-cluster ship only as sdists, so they compile against the
# installed torch — either install from the PyG wheel index or ensure a C++ toolchain:
uv pip install torch-scatter torch-cluster \
  --find-links "https://data.pyg.org/whl/torch-${TORCH_VERSION}+${CUDA}.html"  # e.g. cpu, cu121
uv pip install torchdrug==0.2.1

注意点：

• 先に torch をインストールしてください。TorchDrug と torch-scatter／torch-cluster は、インストール済みの torch に対してグラフ演算をビルドするため、ビルドが一致しないとインポートに失敗します。
• TorchDrug は独自の data.DataLoader、data.Graph、core.Engine を同梱しています。PyTorch 標準のものではなく、こちらを使用してください（以下のループを参照）。
• TorchDrug は rdkit-pypi に依存していますが、2022.09 より新しいバージョンのホイールがありません（Python ≤ 3.11）。これも、スタックが古いインタープリターに固定されている理由の一つです。
• メンテナンスされていません（2023 年 7 月の 0.2.1 以降、リリースがありません）。新しい Python/torch スタックには alterlab-torch-geometric または alterlab-deepchem をお勧めします。TorchDrug のタスク／データセット抽象化（GearNet や逆合成パイプラインなど）が特に必要な場合に、このスキルを使用してください。

クイック例

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

# Load molecular dataset. MoleculeNet-style datasets such as BBBP have no .split()
# method (only datasets with predefined splits do) — split explicitly:
dataset = datasets.BBBP("~/molecule-datasets/")
lengths = [int(0.8 * len(dataset)), int(0.1 * len(dataset))]
lengths += [len(dataset) - sum(lengths)]
train_set, valid_set, test_set = data.ordered_scaffold_split(dataset, lengths)
# (or torch.utils.data.random_split(dataset, lengths) for a random split)

# Define GNN model
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"
)

# Create property prediction task
task = tasks.PropertyPrediction(
    model,
    task=dataset.tasks,
    criterion="bce",
    metric=["auroc", "auprc"]
)

# Train with a native PyTorch loop.
# preprocess() builds the task's MLP prediction head (and label statistics) — core.Engine
# calls it for you; in a manual loop call it BEFORE creating the optimizer, or the head
# is missing and its parameters are never optimized.
task.preprocess(train_set, valid_set, test_set)
# NOTE: use torchdrug.data.DataLoader (NOT torch.utils.data.DataLoader) — its
# default graph_collate packs data.Graph objects; the stock PyTorch collate cannot.
optimizer = torch.optim.Adam(task.parameters(), lr=1e-3)
train_loader = data.DataLoader(train_set, batch_size=32, shuffle=True)

for epoch in range(100):
    for batch in train_loader:
        # task.forward returns (loss, metric), not a bare tensor.
        loss, metric = task(batch)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

標準の（カスタムではない）方法では、core.Engine(task, train_set, valid_set, test_set, optimizer, batch_size=...) がこのループをラップし、バッチ処理とデバイス管理を行います。トレーニングステップを完全に制御する必要がある場合は、上記の手動ループを使用してください。

主な機能

1. 分子特性予測

構造から分子の化学的、物理的、生物学的特性を予測します。

使用例：

• 薬剤らしさと ADMET 特性
• 毒性スクリーニング
• 量子化学特性
• 結合親和性予測

主なコンポーネント：

• 20 以上の分子データセット（BBBP、HIV、Tox21、QM9 など）
• GNN モデル（GIN、GAT、SchNet）
• PropertyPrediction および MultipleBinaryClassification タスク

リファレンス： 詳細については references/molecular_property_prediction.md を参照してください。

• データセットの完全なカタログ
• モデル選択ガイド
• トレーニングのワークフローとベストプラクティス
• 特徴量エンジニアリングの詳細

2. タンパク質モデリング

タンパク質の配列、構造、特性を扱います。

ユースケース：

• 酵素機能の予測
• タンパク質の安定性と溶解性
• 細胞内局在
• タンパク質間相互作用
• 構造予測

主要コンポーネント：

• 15 種類以上のタンパク質データセット（EnzymeCommission、GeneOntology、PDBBind など）
• 配列モデル（ESM、ProteinBERT、ProteinLSTM）
• 構造モデル（GearNet、SchNet）
• さまざまな予測レベルに対応する複数のタスクタイプ

参考資料： 以下については references/protein_modeling.md をご覧ください。

• タンパク質固有のデータセット
• 配列モデルと構造モデル
• 事前学習の戦略
• AlphaFold および ESM との統合

3. ナレッジグラフ推論

生物学的ナレッジグラフ内の欠落したリンクや関係性を予測します。

ユースケース：

• ドラッグリポジショニング
• 疾患メカニズムの発見
• 遺伝子と疾患の関連性
• 複数ホップの生物医学的推論

主要コンポーネント：

• 一般的な KG（FB15k、WN18）と生物医学系（Hetionet）
• 埋め込みモデル（TransE、RotatE、ComplEx）
• KnowledgeGraphCompletion タスク

参考資料： 以下については references/knowledge_graphs.md をご覧ください。

• ナレッジグラフのデータセット（45,000 件の生物医学エンティティを含む Hetionet など）
• 埋め込みモデルの比較
• 評価指標とプロトコル
• 生物医学分野での応用

4. 分子生成

望ましい特性を持つ新規の分子構造を生成します。

ユースケース：

• De novo 創薬
• リード最適化
• 化学空間の探索
• 特性に基づく生成

主要コンポーネント：

• 自己回帰型生成
• GCPN（方策ベースの生成）
• GraphAutoregressiveFlow
• 特性最適化ワークフロー

参考資料： 以下については references/molecular_generation.md をご覧ください。

• 生成戦略（無条件、条件付き、足場ベース）
• 多目的最適化
• 検証とフィルタリング
• 特性予測との統合

5. 逆合成

標的分子から出発物質に至る合成経路を予測します。

ユースケース：

• 合成計画
• 経路の最適化
• 合成容易性の評価
• 多段階計画

主要コンポーネント：

• USPTO-50k 反応データセット
• CenterIdentification（反応中心の予測）
• SynthonCompletion（反応物の予測）
• エンドツーエンドの逆合成パイプライン

参考資料： 以下については references/retrosynthesis.md をご覧ください。

• タスクの分解（center ID → synthon completion）
• 多段階合成計画
• 商業的な入手可能性の確認
• 他の逆合成ツールとの統合

6. グラフニューラルネットワークモデル

さまざまなデータタイプやタスクに対応する GNN アーキテクチャを幅広く紹介します。

利用可能なモデル：

• 汎用 GNN：GCN、GAT、GIN、RGCN、MPNN
• 3D 対応：SchNet、GearNet
• タンパク質固有：ESM、ProteinBERT、GearNet
• ナレッジグラフ：TransE、RotatE、ComplEx、SimplE
• 生成系：GraphAutoregressiveFlow

参考資料： 以下については references/models_architectures.md をご覧ください。

• 各モデルの詳細な説明
• タスクとデータセットに応じたモデル選択ガイド
• アーキテクチャの比較
• 実装のヒント

7. データセット

化学、生物学、ナレッジグラフにまたがる、厳選された 40 種類以上のデータセットです。

カテゴリ：

• 分子特性（創薬、量子化学）
• タンパク質特性（機能、構造、相互作用）
• ナレッジグラフ（一般および生物医学系）
• 逆合成反応

参考資料： 以下については references/datasets.md をご覧ください。

• サイズとタスクを含むデータセットの全一覧
• データセット選択ガイド
• 読み込みと前処理
• 分割戦略（ランダム、足場ベース）

一般的なワークフロー

ワークフロー 1：分子特性の予測

シナリオ： 創薬候補の血液脳関門透過性を予測します。

手順：

1. データセットを読み込みます：datasets.BBBP()
2. モデルを選択します：分子グラフには GIN を使用します
3. タスクを定義します：二値分類を行う PropertyPrediction
4. 現実的な評価のため、足場分割で学習します
5. AUROC と AUPRC を使用して評価します

ナビゲーション： references/molecular_property_prediction.md → データセットの選択 → モデルの選択 → 学習

ワークフロー 2：タンパク質機能の予測

シナリオ： 配列から酵素の機能を予測します。

手順：

1. データセットを読み込みます：datasets.EnzymeCommission()
2. モデルを選択します：ESM（事前学習済み）または GearNet（構造あり）
3. タスクを定義します：多クラス分類を用いた PropertyPrediction
4. 事前学習済みモデルをファインチューニングするか、スクラッチから学習します
5. 正解率とクラスごとの指標を使って評価します

ナビゲーション： references/protein_modeling.md → モデルの選択（配列か構造か）→ 事前学習の戦略

ワークフロー 3：ナレッジグラフを用いたドラッグ・リポジショニング

シナリオ： Hetionet で新しい疾患治療薬を探します。

手順：

1. データセットを読み込みます：datasets.Hetionet()
2. モデルを選択します：RotatE または ComplEx
3. タスクを定義します：KnowledgeGraphCompletion
4. ネガティブサンプリングを用いて学習します
5. 「Compound-treats-Disease」の予測をクエリします
6. 妥当性と作用機序に基づいて絞り込みます

ナビゲーション： references/knowledge_graphs.md → Hetionet データセット → モデルの選択 → バイオメディカルへの応用

ワークフロー 4：分子のデノボ生成

シナリオ： 標的への結合が最適化された、医薬品らしい分子を生成します。

手順：

1. 活性データを用いて特性予測器を学習します
2. 生成手法を選択します：強化学習ベースの最適化には GCPN を使用します
3. 親和性、医薬品らしさ、合成可能性を組み合わせた報酬関数を定義します
4. 特性の制約を設けて候補を生成します
5. 化学的妥当性を検証し、医薬品らしさに基づいて絞り込みます
6. 複数目的のスコアリングに基づいて順位付けします

ナビゲーション： references/molecular_generation.md → 条件付き生成 → 複数目的の最適化

ワークフロー 5：逆合成計画

シナリオ： 標的分子の合成経路を計画します。

手順：

1. データセットを読み込みます：datasets.USPTO50k()
2. 中心同定モデル（RGCN）を学習します
3. シントン補完モデル（GIN）を学習します
4. これらを組み合わせてエンドツーエンドの逆合成パイプラインを構築します
5. 複数ステップの計画には、再帰的に適用します
6. ビルディングブロックの市販状況を確認します

ナビゲーション： references/retrosynthesis.md → タスクの種類 → 複数ステップの計画

統合パターン

RDKit との統合

TorchDrug の分子と RDKit の間で変換します：

from torchdrug import data
from rdkit import Chem

# SMILES → TorchDrug molecule
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, layers
from torchdrug.layers import geometry

# Load AlphaFold predicted structure
protein = data.Protein.from_pdb("AF-P12345-F1-model_v4.pdb")

# Build a residue-level graph with sequential + spatial edges
graph_construction_model = layers.GraphConstruction(
    node_layers=[geometry.AlphaCarbonNode()],
    edge_layers=[
        geometry.SpatialEdge(radius=10.0, min_distance=5),
        geometry.SequentialEdge(max_distance=2),
    ],
    edge_feature="gearnet",
)
# GraphConstruction operates on a packed protein batch, not a bare Protein.
graph = graph_construction_model(data.Protein.pack([protein]))

PyTorch Lightning との統合

Lightning による学習向けにタスクをラップします：

import torch
import pytorch_lightning as pl   # or: import lightning.pytorch as pl

class LightningTask(pl.LightningModule):
    def __init__(self, torchdrug_task):
        super().__init__()
        self.task = torchdrug_task      # call task.preprocess(...) before training

    def training_step(self, batch, batch_idx):
        loss, metric = self.task(batch)  # forward returns (loss, metric)
        return loss

    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)

技術詳細

TorchDrug のアーキテクチャを詳しく知るには、以下を参照してください。

コアコンセプト： references/core_concepts.md には、以下の内容が記載されています。

• アーキテクチャの設計思想（モジュール式、設定可能）
• データ構造（Graph、Molecule、Protein、PackedGraph）
• モデルインターフェースと forward 関数のシグネチャ
• タスクインターフェース（predict、target、forward、evaluate）
• 学習ワークフローとベストプラクティス
• 損失関数と評価指標
• よくある落とし穴とデバッグ

クイックリファレンス早見表

データセットを選択：

• 分子特性 → references/datasets.md → 分子セクション
• タンパク質タスク → references/datasets.md → タンパク質セクション
• ナレッジグラフ → references/datasets.md → ナレッジグラフセクション

モデルを選択：

• 分子 → references/models_architectures.md → GNN セクション → GIN/GAT/SchNet
• タンパク質（配列）→ references/models_architectures.md → タンパク質セクション → ESM
• タンパク質（構造）→ references/models_architectures.md → タンパク質セクション → GearNet
• ナレッジグラフ → references/models_architectures.md → KG セクション → RotatE/ComplEx

一般的なタスク：

• 特性予測 → references/molecular_property_prediction.md または references/protein_modeling.md
• 生成 → references/molecular_generation.md
• 逆合成 → references/retrosynthesis.md
• KG 推論 → references/knowledge_graphs.md

アーキテクチャを理解する：

• データ構造 → references/core_concepts.md → Data Structures
• モデル設計 → references/core_concepts.md → Model Interface
• タスク設計 → references/core_concepts.md → Task Interface

よくある問題のトラブルシューティング

問題：次元不一致エラー → model.input_dim が dataset.node_feature_dim と一致しているか確認してください → references/core_concepts.md → 基本属性を参照してください

問題：分子タスクでのパフォーマンスが低い → ランダム分割ではなく、スキャフォールド分割を使用してください → GCN の代わりに GIN を試してください → references/molecular_property_prediction.md → ベストプラクティスを参照してください

問題：タンパク質モデルが学習しない → 配列タスクには事前学習済みの ESM を使用してください → 構造モデルのエッジ構築を確認してください → references/protein_modeling.md → 学習ワークフローを参照してください

問題：大規模グラフでのメモリエラー → バッチサイズを小さくしてください → 勾配蓄積を使用してください → references/core_concepts.md → メモリ効率を参照してください

問題：生成された分子が無効 → 妥当性制約を追加してください → RDKit による検証で後処理してください → references/molecular_generation.md → 検証とフィルタリングを参照してください

リソース

公式ドキュメント： https://torchdrug.ai/docs/ GitHub： https://github.com/DeepGraphLearning/torchdrug 論文： TorchDrug：創薬のための強力で柔軟な機械学習プラットフォーム

まとめ

タスクに応じて、適切なリファレンスファイルを参照してください。

1. 分子物性予測 → molecular_property_prediction.md
2. タンパク質モデリング → protein_modeling.md
3. ナレッジグラフ → knowledge_graphs.md
4. 分子生成 → molecular_generation.md
5. 逆合成 → retrosynthesis.md
6. モデル選択 → models_architectures.md
7. データセットの選択 → datasets.md
8. 技術的な詳細 → core_concepts.md

各リファレンスでは、例、ベストプラクティス、よくあるユースケースを交えて、それぞれの分野を包括的に解説しています。