torchdrug

TorchDrug

TorchDrug をモジュール式の PyTorch グラフ学習スタックとして使用します。

1. datasets.* データセットを読み込みます。
2. models.* 表現モデルを選択します。
3. tasks.* 目的関数でラップします。
4. core.Engine で学習と評価を行います。

現在の公式ドキュメントと最新リリースは、どちらも 0.2.1 です。 より新しい Python や PyTorch の組み合わせは、互換性を暗黙に仮定せず、未検証として扱ってください。

まずバージョンを確認します

コードを生成またはデバッグする前に、環境を確認します。

python --version
python -c "import torch; print(torch.__version__)"
python -c "import torchdrug; print(torchdrug.__version__)"

TorchDrug 0.2.1 でサポートされる環境は次のとおりです。

• Python 3.7～3.10
• PyTorch 1.8～2.0
• Linux、Windows、macOS
• Apple Silicon：PyTorch 1.13 以降、CPU のみ。MPS はサポートされていません。

プロジェクトで Python 3.11 以降、または PyTorch 2.1 以降を使用する場合は、互換性のある環境を作成するか、ソースからのビルドを明示的にテストしてください。これらの組み合わせをサポート対象として提示しないでください。

インストール

専用の Python 3.10 環境を用意し、TorchDrug のリリースを固定することを推奨します。

uv venv --python 3.10
source .venv/bin/activate
uv pip install "torch==2.0.0"

torch-scatter と torch-cluster のホイールは、PyTorch と CUDA の正確な組み合わせに合わせて、公式インストールページに従ってインストールします。CPU 専用の PyTorch 2.0 環境で再現可能なホイールの組み合わせの一例を示します。

uv pip install "torch-scatter==2.1.1" "torch-cluster==1.6.1" \
  --find-links "https://data.pyg.org/whl/torch-2.0.0+cpu.html"
uv pip install "torchdrug==0.2.1"

CUDA 用ホイールの URL を別の環境に流用しないでください。PyTorch のバージョン、CUDA ビルド、Python ABI、プラットフォームを一致させてください。Apple Silicon では、公式ドキュメントに従い、torch-scatter と torch-cluster をソースからビルドする必要があります。確認済みのソースリビジョンを固定し、CPU で実行されることを想定してください。

標準的なプロパティ予測ワークフロー

ドキュメントに記載されている ClinTox → GIN → PropertyPrediction → Engine のパターンを使用します。

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

dataset = datasets.ClinTox("~/molecule-datasets/")
lengths = [int(0.8 * len(dataset)), int(0.1 * len(dataset))]
lengths.append(len(dataset) - sum(lengths))
train_set, valid_set, test_set = torch.utils.data.random_split(dataset, lengths)

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

optimizer = torch.optim.Adam(task.parameters(), lr=1e-3)
solver = core.Engine(
    task,
    train_set,
    valid_set,
    test_set,
    optimizer,
    batch_size=1024,
)
solver.train(num_epoch=100)
solver.evaluate("valid")

サポート対象の CUDA デバイスが利用できる場合に限り、gpus=[0] を追加します。CPU で実行する場合は、gpus を省略してください。

二値分類の場合、task.predict(batch) はロジットを返します。確率が必要な場合は、torch.sigmoid を適用してください。0.2.1 では、正規化された回帰予測が元のターゲットスケールで返されます。これは旧リリースからの破壊的変更です。

公式ワークフローを選択します

分子プロパティ予測

• データセット：datasets.ClinTox、BBBP、Tox21、QM9、またはその他のドキュメントに記載された分子データセット。
• モデル：まず models.GIN を使用します。選択した特徴量設定でエッジ特徴量が指定される場合は、edge_input_dim を使用します。
• タスク：tasks.PropertyPrediction。
• 分子プロパティ予測を参照してください。

自己教師あり分子事前学習

• InfoGraph：models.InfoGraph(gin_model, separate_model=False) を tasks.Unsupervised でラップします。
• 属性マスキング：tasks.AttributeMasking(model, mask_rate=0.15)。
• ファインチューニング用に同じエンコーダーを再作成し、tasks.PropertyPrediction の学習前に strict=False でチェックポイントを読み込みます。
• 分子プロパティ予測を参照してください。

分子生成

• データセット：datasets.ZINC250k(..., kekulize=True, atom_feature="symbol")。
• GCPN：models.RGCN エンコーダーを tasks.GCPNGeneration でラップします。
• GraphAF：ノードとエッジの models.GraphAF フローを tasks.AutoregressiveGeneration でラップします。
• チュートリアルでサポートされている最適化タスクは "qed" と "plogp" です。基準は "nll" および／または "ppo" です。
• 分子生成を参照してください。

逆合成

• 同期された datasets.USPTO50k のビューを 2 つ作成します。反応モードは反応中心の特定に使用し、as_synthon=True はシントン補完に使用します。
• tasks.CenterIdentification と tasks.SynthonCompletion を個別に学習します。
• 学習済みタスクを tasks.Retrosynthesis と組み合わせます。エンドツーエンドタスクに未加工のモデルを直接渡さないでください。
• 逆合成を参照してください。

ナレッジグラフ推論

• 埋め込みワークフロー：datasets.FB15k237 → models.RotatE → tasks.KnowledgeGraphCompletion。
• ニューラル推論ワークフロー：fact_ratio=0.75 を指定した models.NeuralLP。
• ナレッジグラフ推論を参照してください。

タンパク質モデリング

• data.Protein.from_sequence、from_pdb、または from_molecule を使ってタンパク質を構築します。
• 配列エンコーダーには models.ESM、ProteinCNN、ProteinResNet、 ProteinLSTM、ProteinBERT が含まれます。構造エンコーダーには models.GearNet が含まれます。
• 存在しない便利メソッド protein.residue_graph() ではなく、ドキュメントに記載されたグラフ構築レイヤーを使用します。
• タンパク質モデリングを参照してください。

信頼性の高い TorchDrug コードを書くためのルール

1. 0.2.1 の API に従います。 公式ドキュメントは、常に最新バージョンを反映するサイトではありません。
2. ドキュメントに記載された特徴名を優先します。 atom_feature、bond_feature、residue_feature、mol_feature を使用してください。node_feature、edge_feature、graph_feature は、該当するデータセットコンストラクターでは非推奨のエイリアスです。
3. タスクの前処理は Engine に任せます。 事前学習済みタスクをソルバーを構築せずに組み合わせる場合は、各タスクの preprocess() を手動で呼び出してください。
4. 対応するデータ分割を同期させます。 逆合成では、反応データセットとシントンデータセットを分割する前に、同じ乱数シードを再設定してください。
5. TorchDrug のコレーションを使用します。 data.graph_collate または core.Engine を使用してください。一般的な PyTorch のコレーションでは、TorchDrug のグラフをパックする方法がわかりません。
6. モデル、タスク、エンジンの引数を分けます。 コードを誤って作成するよくある原因は、タスクのオプションをモデルに渡すことや、合成されたタスクが必要な場所に生のモデルを渡すことです。
7. 生成された化学構造を検証します。 モデルの出力は候補として扱い、実験的に有効である、または合成可能であるとは見なさないでください。

トラブルシューティング

インストールまたはインポートの失敗

Python、PyTorch、torch-scatter、torch-cluster の互換性をまとめて確認してください。ほとんどの失敗は、バイナリホイールの不一致、サポート対象外の Python バージョン、または MPS を使用しようとしたことが原因です。

特徴量の次元の不一致

読み込んだデータセットからモデルの次元を取得してください。

• dataset.node_feature_dim
• dataset.edge_feature_dim
• dataset.num_bond_type
• ナレッジグラフの場合は dataset.num_entity と dataset.num_relation

異なる特徴量設定からコピーした次元をハードコードしないでください。

デバイスの不一致

サポート対象の CUDA 実行では、core.Engine に gpus=[0] を渡してください。手動で予測する場合は、まずコレーションを行ってから、ネストされたバッチ全体を utils.cuda で移動してください。

チェックポイントの不一致

同じモデルと特徴量設定を再作成してください。事前学習からファインチューニングへ転移する場合は、チェックポイントの "model" の状態を strict=False で読み込みます。ソルバー全体の場合は、solver.save() と solver.load() を使用してください。

リファレンス一覧

• コアコンセプトとデータ構造
• データセット
• モデルとアーキテクチャ
• 分子特性予測と事前学習
• タンパク質モデリング
• 分子生成
• 逆合成
• ナレッジグラフ推論

上流のソース

• TorchDrug 0.2.1 ドキュメント
• チュートリアル一覧
• インストール
• パッケージリファレンス
• TorchDrug 0.2.1 リリースノート

エージェントの作業手順

1. 環境を確認します。 ツールキットのバージョンと入力形式（SMILES、SDF、PDB、CIF）を確認し、計算前に分子をサニタイズします。
2. 入力を明確にします。 データまたはユーザーから、形式、識別子、パラメーターを確認します。結果を左右する値は推測せず、質問します。
3. まず小規模に実行します。 少数の分子、または1つのターゲットを処理し、出力を確認します。
4. 上記の手順と参照情報に従って、タスク全体を実行します。
5. 結果を検証します。 無効な構造、単位の一貫性、予測モデルの適用範囲を確認します。
6. 報告します。 実行した内容（バージョン、コマンド、パラメーター）、確認した内容、未確定の点を明記します。

このような場合	次の対応
分子を解析またはサニタイズできない	理由を添えて報告し、対象をスキップします。構造を黙って除外したり変更したりしてはいけません。
これらの手順に記載された関数、フラグ、エンドポイントがインストール済みバージョンに存在しない	インストール済みバージョンのドキュメント（help()、--help、公式ドキュメント）を確認し、適宜対応してユーザーに伝えます。APIを捏造してはいけません。
必須の入力、識別子、パラメーターが曖昧	ユーザーに質問するか、実行前に仮定を明示します。

整合性に関するルール

• 結果、パラメーター、識別子、引用、統計情報を捏造してはいけません。実行または検証できない場合は、その旨を明確に伝えます。
• 計算による予測は仮説であり、実験結果ではありません。そのように明記します。
• ここに記載されたバージョン固有の情報は古くなっている可能性があります。インストール済みバージョンの公式ドキュメントで確認します。
• 費用が発生する操作、共有 GPU やクラウドのクォータを消費する操作、個人データや患者データを扱う操作、または元に戻せない操作の前には、確認を求めます。

関連スキル

• deepchem：多様な特徴量化手法と事前構築済みデータセットを使った分子機械学習。
• neurokit2：NeuroKit2 を使用して、生理時系列の前処理やイベント／間隔分析などのための再現可能な研究ワークフローを構築または監査します。…
• rowan：Rowan は、Python API を備えたクラウドネイティブの分子モデリングおよびメディシナルケミストリーのワークフロープラットフォームです。