MACEは、高次の等しいメッセージの合格を伴う高速かつ正確な機械学習間ポテンシャルを提供します。
このリポジトリには、Ilyes Batatia、Gregor Simm、David Kovacs、Gabor Csanyiのグループ、および友人が開発したMACEリファレンス実装が含まれています(寄稿者を参照)。
また利用可能:
部分的なドキュメントは、https://mace-docs.readthedocs.ioで入手できます
必ずPytorchをインストールしてください。インストール手順については、公式のPytorchインストールを参照してください。システムに適切なオプションを選択します。
これは、MACEをインストールする推奨方法です。
pip install --upgrade pip
pip install mace-torch
注: Pypiの同名のパッケージは、これとは何の関係もありません。
git clone https://github.com/ACEsuit/mace.git
pip install ./mace
MACEモデルをトレーニングするには、 mace_run_train
スクリプトを使用できます。これは、PIPがバイナリを配置する通常の場所にある必要があります(または、 python3 <path_to_cloned_dir>/mace/cli/run_train.py
明示的に実行できます)
mace_run_train
--name= " MACE_model "
--train_file= " train.xyz "
--valid_fraction=0.05
--test_file= " test.xyz "
--config_type_weights= ' {"Default":1.0} '
--E0s= ' {1:-13.663181292231226, 6:-1029.2809654211628, 7:-1484.1187695035828, 8:-2042.0330099956639} '
--model= " MACE "
--hidden_irreps= ' 128x0e + 128x1o '
--r_max=5.0
--batch_size=10
--max_num_epochs=1500
--swa
--start_swa=1200
--ema
--ema_decay=0.99
--amsgrad
--restart_latest
--device=cuda
特定の検証セットを提供するには、引数--valid_file
を使用します。検証セットを評価するために大きなバッチサイズを設定するには、 --valid_batch_size
を指定します。
モデルのサイズを制御するには、 --hidden_irreps
変更する必要があります。ほとんどのアプリケーションでは、推奨されるデフォルトのモデルサイズは、 --hidden_irreps='256x0e'
(256の不変メッセージ)または--hidden_irreps='128x0e + 128x1o'
です。モデルが十分に正確でない場合は、例えば128x0e + 128x1o + 128x2e
などの高次機能を含めるか、チャネルの数を256
に増やすことができます。 --num_channels=128
および--max_L=1
キーを使用してモデルを指定することもできます。
通常、上記の例のようにコマンドラインを通してエネルギーを読むよりも、孤立した原子をトレーニングセットに追加することが好まれます。トレーニングセットにラベルを付けるには、情報フィールドにconfig_type=IsolatedAtom
を設定します。分離された原子のエネルギーを使用しないか、知らない場合は、最小二乗回帰を使用して原子エネルギーを推定するオプション--E0s="average"
を使用できます。
キーワード--swa
が有効になっている場合、トレーニングエポックの最後の約20%( --start_swa
エポック)の損失のエネルギー重量が増加します。この設定は通常、エネルギーエラーを下げるのに役立ちます。
キーワード--default_dtype
を使用して精度を変更できます。デフォルトはfloat64
ですが、 float32
大幅なスピードアップ(通常はトレーニングでx2の係数です)を提供します。
キーワード--batch_size
および--max_num_epochs
、トレーニングセットのサイズに基づいて調整する必要があります。トレーニングデータの数が増加すると、バッチサイズを増やし、エポックの数を減らす必要があります。初期設定のヒューリスティックは、勾配更新定数の数を200 000に考慮することです。
このコードは、ストレスのあるバルク構造と分離分子の両方を含む、不均一なラベルを使用したトレーニングセットを処理できます。この例では、コードに分子に対するストレスを無視させるには、分子構成A config_stress_weight = 0.0
に追加されます。
Apple Silicon GPUアクセラレーションを使用するには、最新のPytorchバージョンをインストールし、 --device=mps
を指定してください。
マルチGPUトレーニングには、 --distributed
フラグを使用します。これにより、PytorchのDistributedDatapar Allel Moduleを使用して、複数のGPUでモデルをトレーニングします。大きなデータセットのオンラインデータロードと組み合わせます(以下を参照)。 SluRMスクリプトの例はmace/scripts/distributed_example.sbatch
にあります。
YAMLを使用してすべてまたは一部の引数を解析するオプションが利用可能です。たとえば、上記の引数を使用してモデルをトレーニングするには、次のコンテンツでYAMLファイルyour_configs.yaml
を作成できます。
name : nacl
seed : 2024
train_file : train.xyz
swa : yes
start_swa : 1200
max_num_epochs : 1500
device : cpu
test_file : test.xyz
E0s :
41 : -1029.2809654211628
38 : -1484.1187695035828
8 : -2042.0330099956639
config_type_weights :
Default : 1.0
コマンドライン--config="your_configs.yaml"
に追加されます。コマンドラインで指定された引数は、YAMLファイルにある引数を上書きします。
XYZファイルでMACEモデルを評価するには、 mace_eval_configs
を実行します。
mace_eval_configs
--configs= " your_configs.xyz "
--model= " your_model.model "
--output= " ./your_output.xyz "
Colabチュートリアルを実行して、メイスをすばやく開始できます。
また、より詳細なコラブチュートリアルもあります。
GPUメモリに適合しない可能性のある大きなデータセットがある場合は、CPU上のデータを前処理し、モデルをトレーニングするためにオンラインデータロードを使用することをお勧めします。 XYZファイルとして指定されたデータセットを事前に処理するには、 preprocess_data.py
スクリプトを実行します。例はここに記載されています:
mkdir processed_data
python ./mace/scripts/preprocess_data.py
--train_file= " /path/to/train_large.xyz "
--valid_fraction=0.05
--test_file= " /path/to/test_large.xyz "
--atomic_numbers= " [1, 6, 7, 8, 9, 15, 16, 17, 35, 53] "
--r_max=4.5
--h5_prefix= " processed_data/ "
--compute_statistics
--E0s= " average "
--seed=123
すべてのオプションとそれらの少しの説明を表示するにはpython ./mace/scripts/preprocess_data.py --help
を実行します。スクリプトは、トレーニングに使用できるprocessed_data
フォルダーに多数のHDF5ファイルを作成します。トレーニング用の1つのフォルダーがあります。1つは検証用と、テストセットの各config_type
用の別のフォルダーがあります。モデルをトレーニングするには、次のようにrun_train.py
スクリプトを使用します。
python ./mace/scripts/run_train.py
--name= " MACE_on_big_data "
--num_workers=16
--train_file= " ./processed_data/train.h5 "
--valid_file= " ./processed_data/valid.h5 "
--test_dir= " ./processed_data "
--statistics_file= " ./processed_data/statistics.json "
--model= " ScaleShiftMACE "
--num_interactions=2
--num_channels=128
--max_L=1
--correlation=3
--batch_size=32
--valid_batch_size=32
--max_num_epochs=100
--swa
--start_swa=60
--ema
--ema_decay=0.99
--amsgrad
--error_table= ' PerAtomMAE '
--device=cuda
--seed=123
重みとバイアスを備えたメイスを使用して実験をログインしたい場合は、単にインストールしてください
pip install ./mace[wandb]
必要なキーワード引数( --wandb
、 --wandb_project
、 --wandb_entity
、 --wandb_name
、 --wandb_log_hypers
)を指定します。
マテリアルプロジェクト(MP)と協力して、MP緩和軌道から選択されたMPTRJデータセットの1.6 mバルク結晶の89の要素をカバーするユニバーサルメイスポテンシャルをトレーニングしました。モデルはhttps://github.com/acesuit/mace-mpのGithubで解放されます。それらを使用する場合は、幅広い範囲のアプリケーションとベンチマークを含む私たちの論文を引用してください。
注意
MACE-MPモデルは、VASP出力からのMPTRJ RAW DFTエネルギーでトレーニングされており、MPのDFTエネルギーまたはCHGNETのエネルギーに直接匹敵しません。補正)、および14個のアニオン種(アニオン補正)。詳細については、MPドキュメントとMP2020Compatibility.YAMLを参照してください。
from mace . calculators import mace_mp
from ase import build
atoms = build . molecule ( 'H2O' )
calc = mace_mp ( model = "medium" , dispersion = False , default_dtype = "float32" , device = 'cuda' )
atoms . calc = calc
print ( atoms . get_potential_energy ())
シリーズ(小、中、大)転送可能な有機力場があります。これらは、有機分子、結晶、分子液のシミュレーションに、または新しいデータセットでの微調整の出発点として使用できます。モデルはASLライセンスの下でリリースされます。モデルはhttps://github.com/acesuit/mace-offのGithubで解放されます。それらを使用する場合は、詳細なベンチマークとサンプルアプリケーションも含まれている私たちの論文を引用してください。
from mace . calculators import mace_off
from ase import build
atoms = build . molecule ( 'H2O' )
calc = mace_off ( model = "medium" , device = 'cuda' )
atoms . calc = calc
print ( atoms . get_potential_energy ())
MACE-MP-0ファンデーションモデルのいずれかをFinetuneするには、 mace_run_train
スクリプトを追加引数--foundation_model=model_type
で使用できます。たとえば、新しいデータセットで小さなモデルをFintuneするには、以下を使用できます。
mace_run_train
--name= " MACE "
--foundation_model= " small "
--train_file= " train.xyz "
--valid_fraction=0.05
--test_file= " test.xyz "
--energy_weight=1.0
--forces_weight=1.0
--E0s= " average "
--lr=0.01
--scaling= " rms_forces_scaling "
--batch_size=2
--max_num_epochs=6
--ema
--ema_decay=0.99
--amsgrad
--default_dtype= " float32 "
--device=cuda
--seed=3
その他のオプションは、「中」および「大規模」、またはファンデーションモデルへのパスです。別のモデルをFintuneする場合、モデルは提供されたパス--foundation_model=$path_model
からロードされますが、モデルに一致するハイパーパラメーター(非表示のIRREPS、R_MAXなど)の完全なセットを提供する必要があります。
このプロジェクトでは、事前コミットを使用して、CODEのコードフォーマットとリントを実行します。また、 black
、 isort
、 pylint
、 mypy
も使用しています。開発環境をPython環境にインストールして、 dev
環境をセットアップすることをお勧めします。
pip install -e " .[dev] "
pre-commit install
2行目は、CommitのAutomaticaly Run Code Checksにpre-commit
を初期化します。これをチェックするためにCIを設定していますが、誤って悪いコードをコミットしないようにする(およびプッシュ)する前にそれらのコマンドを実行することを強くお勧めします。
MITライセンスの下でプルリクエストを受け入れて喜んでいます。コメントとしてライセンステキストをプルリクエストにコピー/貼り付けてください。
このコードを使用する場合は、私たちの論文を引用してください。
@inproceedings { Batatia2022mace ,
title = { {MACE}: Higher Order Equivariant Message Passing Neural Networks for Fast and Accurate Force Fields } ,
author = { Ilyes Batatia and David Peter Kovacs and Gregor N. C. Simm and Christoph Ortner and Gabor Csanyi } ,
booktitle = { Advances in Neural Information Processing Systems } ,
editor = { Alice H. Oh and Alekh Agarwal and Danielle Belgrave and Kyunghyun Cho } ,
year = { 2022 } ,
url = { https://openreview.net/forum?id=YPpSngE-ZU }
}
@misc { Batatia2022Design ,
title = { The Design Space of E(3)-Equivariant Atom-Centered Interatomic Potentials } ,
author = { Batatia, Ilyes and Batzner, Simon and Kov{'a}cs, D{'a}vid P{'e}ter and Musaelian, Albert and Simm, Gregor N. C. and Drautz, Ralf and Ortner, Christoph and Kozinsky, Boris and Cs{'a}nyi, G{'a}bor } ,
year = { 2022 } ,
number = { arXiv:2205.06643 } ,
eprint = { 2205.06643 } ,
eprinttype = { arxiv } ,
doi = { 10.48550/arXiv.2205.06643 } ,
archiveprefix = { arXiv }
}
ご質問がある場合は、[email protected]までお問い合わせください。
バグまたは機能リクエストについては、GitHubの問題を使用してください。
MACEは、MITライセンスの下で公開および配布されます。