概要
TorchMD-Netは、単一の学習済み予測モデルではなく、分子および生体分子系のニューラルネットワークポテンシャルを学習・利用するためのツールキットです。ポテンシャルはPyTorchモジュールとして提供され、READMEには分子動力学ワークフロー向けのACEMD、OpenMM、TorchMDとの連携が記載されています。利用可能なアーキテクチャとしてTensorNetとTensorNet2が挙げられており、Equivariant Transformer、Transformer、Graph Neural Networkの各アーキテクチャは非推奨とされています。
学習はYAMLファイルまたはコマンドライン引数で設定でき、コマンドラインの値が優先されます。自作のNumPyデータセットについて、ドキュメントで指定されている入力は原子種と原子座標で、それらにエネルギーラベル、力ラベル、または両方を対応付けます。別のデータセット実装では、zとposに加え、y、neg_dy、または両方を含むtorch-geometric Dataオブジェクトを返します。これらのインターフェースは、ラベル付き構造から分子エネルギーと座標微分を学習するワークフローを支援します。リポジトリには、自作priorモデルや新しいアーキテクチャの拡張方法も記載されています。
推論向けに、TorchMD-Netはチェックポイントの読み込みとASE計算器インターフェースを提供しており、学習済みAceFFモデルの例も含まれています。ローダーは、ドキュメントに記載された旧形式のマーカーを検出すると、旧形式のTensorNetおよびTensorNet2チェックポイントのレイアウトを自動的に再マッピングします。マルチノード学習に対応していますが、READMEでは各ノードのGPU数を同じにする必要があるとされており、GPUアーキテクチャが異なる場合の注意点や、複数GPUでハングする可能性について警告しています。出典の抜粋からは、特定の化学分野における予測精度は確認できません。リポジトリのMITコードライセンスを、別途配布されるモデルの重みやデータセットにも同じ条件が適用される根拠として扱うべきではありません。
主な機能
- TensorNetおよびTensorNet2のニューラルネットワークポテンシャルをPyTorchモジュールとして実装。ET、T、GNアーキテクチャは非推奨と記載。
- YAMLファイルとコマンドライン引数による学習設定に対応し、コマンドライン引数で上書き可能。
- 原子種、座標、エネルギーまたは力のラベルを含む自作NumPyデータセットを読み込み、torch-geometricのデータセット実装にも対応。
- `load_model`による学習済みチェックポイントの読み込みと、`TMDNETCalculator`によるASEインターフェースを提供。AceFFの例も掲載。
- ドキュメントに記載された旧形式のTensorNetおよびTensorNet2チェックポイントのテンソル配置を自動検出して再マッピングし、明示的な上書きオプションも提供。
- 自作priorモデル、アーキテクチャの拡張、GPU数に制約のあるマルチノード学習について説明。
用途
- 評価案：ラベル付き構造の集合で分子ポテンシャルを学習し、保留した構造に対するエネルギーおよび力の予測を評価する。
- 評価案：ドキュメントに記載されたASE計算器ワークフローを通じてAceFFチェックポイントを読み込み、対象の分子系への適合性を評価する。
- 評価案：学習済みポテンシャルをACEMD、OpenMM、またはTorchMDのシミュレーションワークフローで使用する可能性を調べ、関連する連携要件を確認する。
- 評価案：リポジトリの拡張ガイドを使って自作priorモデルまたはアーキテクチャを試作し、適切な検証データセットで評価する。
使い方
- ドキュメントとリポジトリのREADMEを読みます。ポテンシャルを学習するか、既存のチェックポイントを読み込むかを選び、非推奨のアーキテクチャを確認します。
- READMEに記載されたpipまたはconda/mambaによるインストール手順に従います。例にあるCUDAの選択が自分の環境に適していると決めつけず、リンク先のPyTorchの手順を参考に、PyTorchのアクセラレーター対応ビルドを選択します。
- 学習する場合は、ドキュメントに記載された自作データセットインターフェースを使い、原子種、座標、エネルギーまたは力のラベルを用意します。設定例を参照し、コマンドラインの設定がYAMLの値より優先されることに注意してください。
- 学習済みモデルで推論する場合は、AceFFの例またはチェックポイント読み込みの説明を参照します。自動検出を上書きする前に、ドキュメントに記載された旧形式レイアウトへの対応を確認してください。
- 評価案として、シミュレーションに使用する前に、代表的な保留構造に対する予測を評価します。マルチノード学習へ拡張する場合は、READMEにある各ノードのGPU数を同じにする要件に従い、ハードウェアと通信に関する制約を確認してください。