この作品の詳細については当社ウェブサイトをご覧ください。
このリポジトリには、論文「分子から材料へ: 原子特性予測のための大規模一般化可能モデルの事前トレーニング」のコードが含まれています。
このリポジトリは非推奨となり、積極的にメンテナンスされなくなりました。現在、このリポジトリの機能を Open Catalyst プロジェクトの公式リポジトリに統合する作業を行っています。このリポジトリに関してご質問や懸念がある場合は、お気軽に github の問題を作成するか、電子メールでご連絡ください。 Nima Shoghi と CC Brandon Wood に電子メールを送信してください。
この研究では、ジョイント マルチドメイン事前トレーニング (JMP) を導入します。これは、異なる化学ドメインからの複数のデータセットを同時にトレーニングし、各データセットをマルチタスク フレームワーク内の固有の事前トレーニング タスクとして扱う教師あり事前トレーニング戦略です。 。結合されたトレーニング データセットは、OC20、OC22、ANI-1x、Transition-1x の約 1 億 2,000 万のシステムで構成されています。
この作業の主な貢献は次のとおりです。
私たちは、小分子、大分子、材料にわたる多様なベンチマークスイート全体での微調整パフォーマンスを評価することにより、JMP の強力な一般化能力を実証します。 JMP は一貫してゼロからのトレーニングよりも優れたパフォーマンスを示し、40 の微調整ベンチマークのうち 34 で最先端のベンチマークを設定または一致させます。
JMP を使用すると、小さなデータセットで最初からトレーニングすると通常はオーバーフィットしてしまう、より大きなモデルへの効率的なスケーリングが可能になることを示します。事前トレーニングは強力な正則化機能として機能し、複数の低データ ベンチマークで新しい最先端のパフォーマンスを設定する 2 億 3,500 万のパラメーターを使用してモデルをトレーニングできるようになります。
JMP の計算要件を詳細に分析します。初期費用はかかりますが、JMP の事前トレーニング コストは、最初からトレーニングする場合と比較して 12 倍以上高速な微調整を可能にすることで回収できることを示しています。
多様な化学データで大規模なモデルを事前トレーニングすることにより、JMP は化学に対する普遍的な ML の可能性という目標に向けた重要な一歩となると私たちは考えています。利用可能なデータと計算能力が継続的に増加すると、転送可能なアトミック表現を学習する JMP の能力が向上するだけです。
JMP は、ゼロからトレーニングする場合と比べて平均 59% の向上を示し、40 タスク中 34 タスクで最先端のレベルに適合または設定します。私たちの研究は、特にデータ量の少ないタスクにおいて、化学ドメイン全体の特性予測を進めるために多様なデータを利用する事前トレーニング戦略の可能性を浮き彫りにしています。
まず、リポジトリのクローンを作成し、ルート ディレクトリに移動します。
git clone https://github.com/facebookresearch/JMP.git
cd JMP次に、 environment.ymlファイルで提供されているように、conda 環境をセットアップします。これを行うには、次のコマンドを実行します (注意: conda がインストールされている場合は、 conda mambaに置き換えます)。
conda env create -f environment.yml -n jmp上記のコマンドが失敗した場合は、自分で環境を手動で作成できます。
# Create the environment
conda create -n jmp python=3.11
conda activate jmp
# Install PyTorch
conda install -y -c pytorch -c nvidia pytorch torchvision torchaudio pytorch-cuda=12.1
# Install PyG, PyTorch Scatter, PyTorch Sparse.
conda install -c pyg pyg pytorch-sparse pytorch-cluster
# Install other conda dependencies
conda install -y
-c conda-forge
numpy matplotlib seaborn sympy pandas numba scikit-learn plotly nbformat ipykernel ipywidgets tqdm pyyaml networkx
pytorch-lightning torchmetrics lightning
einops wandb
cloudpickle
" pydantic>2 "
frozendict wrapt varname typing-extensions lovely-tensors lovely-numpy requests pytest nbval
# Install pip dependencies
pip install lmdb
# Install dependencies for materials datasets
pip install ase
# Install dependencies for large molecule datasets
conda install h5py
# Install MatBench dependencies
pip install matbench
# Install dependencies for PDBBind
pip install biopython rdkit
# Install dependencies for pre-processing ANI1x/Transition1x
pip install multiprocess次に、環境をアクティブ化します。
conda activate jmp最後に、次のように現在のパッケージをインストールします。
pip install -e .これでコードを使用する準備が整いました。コードの実行方法の例については、 configsディレクトリ内の構成ファイルを参照してください。
事前トレーニングと微調整に使用されるデータは、サイズの制約によりこのリポジトリには含まれません。ただし、OC20、OC22、ANI-1x、Transition-1x、QM9、rMD17、MatBench、QMOF、SPICE、および MD22 データセットのダウンロードと前処理の手順を以下に示します。
ANI-1x データセットをダウンロードし、コードベースで使用できる形式に変換するには、次の手順に従ってください。
mkdir -p /path/to/datasets/ani1x
cd /path/to/datasets/ani1x.h5形式) をダウンロードします。 wget https://springernature.figshare.com/ndownloader/files/18112775 -O ani1x-release.h5python -m jmp.datasets.scripts.ani1x_preprocess.ani1x_splits --input_file ani1x-release.h5 --train_keys_output train_keys.pkl --val_keys_output val_keys.pkl --test_keys_output test_keys.pkl.trajファイルのセットに変換します ( --num_workers使用可能な CPU コアの数に設定します)。 mkdir -p traj
mkdir -p traj/train traj/val traj/test
python -m jmp.datasets.scripts.ani1x_preprocess.ani1x_write_traj --ani1x_h5 ani1x-release.h5 --split_keys train_keys.pkl --split train --traj_dir traj/train --num_workers 32
python -m jmp.datasets.scripts.ani1x_preprocess.ani1x_write_traj --ani1x_h5 ani1x-release.h5 --split_keys val_keys.pkl --split val --traj_dir traj/val --num_workers 32
python -m jmp.datasets.scripts.ani1x_preprocess.ani1x_write_traj --ani1x_h5 ani1x-release.h5 --split_keys test_keys.pkl --split test --traj_dir traj/test --num_workers 32.trajファイルを.lmdbファイルに変換します。 mkdir -p lmdb
mkdir -p lmdb/train lmdb/val lmdb/test
python -m jmp.datasets.scripts.ani1x_preprocess.ani1x_write_lmdbs --data_path traj/train --out_path lmdb/train --split train --num_workers 32
python -m jmp.datasets.scripts.ani1x_preprocess.ani1x_write_lmdbs --data_path traj/val --out_path lmdb/val --split val --num_workers 32
python -m jmp.datasets.scripts.ani1x_preprocess.ani1x_write_lmdbs --data_path traj/test --out_path lmdb/test --split test --num_workers 32python -m jmp.datasets.scripts.ani1x_preprocess.ani1x_linear_ref linref --src lmdb/train --out_path linref.npzpython -m jmp.datasets.scripts.ani1x_preprocess.ani1x_linear_ref compute_mean_std --src lmdb/train --linref_path linref.npz --out_path mean_std.pklTransition-1x データセットをダウンロードし、コードベースで使用できる形式に変換するには、次の手順に従ってください。
mkdir -p /path/to/datasets/transition1x
cd /path/to/datasets/transition1x.h5形式) をダウンロードします。 wget https://figshare.com/ndownloader/files/36035789 -O transition1x-release.h5.trajファイルのセットに変換します ( --num_workers使用可能な CPU コアの数に設定します)。 mkdir -p traj
mkdir -p traj/train traj/val traj/test
python -m jmp.datasets.scripts.transition1x_preprocess.trans1x_write_traj --transition1x_h5 transition1x-release.h5 --split train --traj_dir traj/train --num_workers 32
python -m jmp.datasets.scripts.transition1x_preprocess.trans1x_write_traj --transition1x_h5 transition1x-release.h5 --split val --traj_dir traj/val --num_workers 32
python -m jmp.datasets.scripts.transition1x_preprocess.trans1x_write_traj --transition1x_h5 transition1x-release.h5 --split test --traj_dir traj/test --num_workers 32.trajファイルを.lmdbファイルに変換します。 mkdir -p lmdb
mkdir -p lmdb/train lmdb/val lmdb/test
python -m jmp.datasets.scripts.transition1x_preprocess.trans1x_write_lmdbs --data_path traj/train --out_path lmdb/train --split train --num_workers 32
python -m jmp.datasets.scripts.transition1x_preprocess.trans1x_write_lmdbs --data_path traj/val --out_path lmdb/val --split val --num_workers 32
python -m jmp.datasets.scripts.transition1x_preprocess.trans1x_write_lmdbs --data_path traj/test --out_path lmdb/test --split test --num_workers 32python -m jmp.datasets.scripts.transition1x_preprocess.trans1x_linear_ref linref --src lmdb/train --out_path linref.npzpython -m jmp.datasets.scripts.transition1x_preprocess.trans1x_linear_ref compute_mean_std --src lmdb/train --linref_path linref.npz --out_path mean_std.pklpython -m jmp.datasets.finetune.rmd17 download --destination /path/to/datasets/rmd17/python -m jmp.datasets.finetune.qm9 download --destination /path/to/datasets/qm9/python -m jmp.datasets.finetune.md22 download --destination /path/to/datasets/md22/python -m jmp.datasets.finetune.spice download --destination /path/to/datasets/spice/python -m jmp.datasets.finetune.mat_bench download --destination /path/to/datasets/matbench/python -m jmp.datasets.finetune.qmof download --destination /path/to/datasets/qmof/ 事前トレーニングされたチェックポイントは、次のリンクからダウンロードできます。
私たちのコードベースは Python ベースになるように設計されています。モデルをトレーニングするには、モデル、データセット、トレーニング パラメーターを指定する構成ファイルを作成します。これらの構成オブジェクトは、Pydantic を使用して完全に型チェックおよび検証されます。
設定オブジェクトを作成したら、 jmp.lightning.Runnerクラスを使用してトレーニング ループをラップできます。モデルを微調整する方法の例については、 configs/jmp_l_finetune.ipynbを参照してください。
LICENSEファイルにあるように、JMP の大部分は CC-BY-NC ライセンスを取得しています。ただし、プロジェクトの一部は別のライセンス条項に基づいて利用可能です。
研究でこのコードを使用する場合は、次の論文を引用してください。
@article{shoghi2023molecules,
title={From molecules to materials: Pre-training large generalizable models for atomic property prediction},
author={Shoghi, Nima and Kolluru, Adeesh and Kitchin, John R and Ulissi, Zachary W and Zitnick, C Lawrence and Wood, Brandon M},
journal={arXiv preprint arXiv:2310.16802},
year={2023}
}
