mace
v0.3.8
MACE menyediakan potensi interatomik pembelajaran mesin yang cepat dan akurat dengan lewat pesan equivariant yang lebih tinggi.
Repositori ini berisi implementasi referensi MACE yang dikembangkan oleh Ilyes Batatia, Gregor Simm, David Kovacs, dan kelompok Gabor Csanyi, dan teman -teman (lihat Kontributor).
Juga tersedia:
Dokumentasi parsial tersedia di: https://mace-docs.readthedocs.io
Pastikan untuk menginstal Pytorch. Silakan merujuk ke instalasi Pytorch resmi untuk instruksi instalasi. Pilih opsi yang sesuai untuk sistem Anda.
Ini adalah cara yang disarankan untuk menginstal MACE.
pip install --upgrade pip
pip install mace-torchCatatan: Paket homonim di PYPI tidak ada hubungannya dengan yang satu ini.
git clone https://github.com/ACEsuit/mace.git
pip install ./mace Untuk melatih model tongkat, Anda dapat menggunakan skrip mace_run_train , yang seharusnya berada di tempat yang biasa tempat Pip menempatkan binari (atau Anda dapat secara eksplisit menjalankan 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 Untuk memberikan set validasi tertentu, gunakan argumen --valid_file . Untuk mengatur ukuran batch yang lebih besar untuk mengevaluasi set validasi, tentukan --valid_batch_size .
Untuk mengontrol ukuran model, Anda perlu mengubah --hidden_irreps . Untuk sebagian besar aplikasi, ukuran model default yang disarankan adalah --hidden_irreps='256x0e' (artinya 256 pesan invarian) atau --hidden_irreps='128x0e + 128x1o' . Jika model tidak cukup akurat, Anda dapat memasukkan fitur orde lebih tinggi, misalnya, 128x0e + 128x1o + 128x2e , atau meningkatkan jumlah saluran menjadi 256 . Dimungkinkan juga untuk menentukan model menggunakan --num_channels=128 dan --max_L=1 tombol.
Biasanya lebih disukai untuk menambahkan atom yang terisolasi ke set pelatihan, daripada membaca dalam energi mereka melalui baris perintah seperti pada contoh di atas. Untuk memberi label pada set pelatihan, atur config_type=IsolatedAtom di bidang info mereka. Jika Anda lebih suka tidak menggunakan atau tidak mengetahui energi atom yang terisolasi, Anda dapat menggunakan opsi --E0s="average" yang memperkirakan energi atom menggunakan regresi kuadrat terkecil.
Jika kata kunci --swa diaktifkan, berat energi kerugian meningkat untuk ~ 20% terakhir dari zaman pelatihan (dari --start_swa zaman). Pengaturan ini biasanya membantu menurunkan kesalahan energi.
Presisi dapat diubah menggunakan kata kunci --default_dtype , standarnya float64 tetapi float32 memberikan percepatan yang signifikan (biasanya faktor X2 dalam pelatihan).
Kata kunci --batch_size dan --max_num_epochs harus diadaptasi berdasarkan ukuran set pelatihan. Ukuran batch harus ditingkatkan ketika jumlah data pelatihan meningkat, dan jumlah zaman harus dikurangi. Heuristik untuk pengaturan awal, adalah mempertimbangkan jumlah pembaruan gradien yang konstan hingga 200.000, yang dapat dihitung sebagai
Kode ini dapat menangani set pelatihan dengan label heterogen, misalnya mengandung kedua struktur curah dengan stres dan molekul terisolasi. Dalam contoh ini, untuk membuat kode mengabaikan tekanan pada molekul, tambahkan konfigurasi molekul Anda A config_stress_weight = 0.0 .
Untuk menggunakan akselerasi GPU Apple Silicon, pastikan untuk menginstal versi Pytorch terbaru dan tentukan --device=mps .
Untuk pelatihan multi-GPU, gunakan bendera --distributed . Ini akan menggunakan modul DistributedDataParallel Pytorch untuk melatih model pada beberapa GPU. Kombinasikan dengan pemuatan data on-line untuk kumpulan data besar (lihat di bawah). Contoh skrip slurm dapat ditemukan di mace/scripts/distributed_example.sbatch .
Opsi untuk menguraikan semua atau beberapa argumen menggunakan YAML tersedia. Misalnya, untuk melatih model menggunakan argumen di atas, Anda dapat membuat file yaml yaml your_configs.yaml dengan konten berikut:
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
Dan tambahkan ke baris perintah --config="your_configs.yaml" . Argumen apa pun yang ditentukan dalam baris perintah akan menimpa yang ada di file YAML.
Untuk mengevaluasi model MACE Anda pada file XYZ, jalankan mace_eval_configs :
mace_eval_configs
--configs= " your_configs.xyz "
--model= " your_model.model "
--output= " ./your_output.xyz " Anda dapat menjalankan tutorial Colab kami untuk memulai dengan cepat dengan Mace.
Kami juga memiliki tutorial Colab yang lebih rinci di:
Jika Anda memiliki dataset besar yang mungkin tidak masuk ke dalam memori GPU, disarankan untuk preprocess data pada CPU dan menggunakan dataloading on-line untuk melatih model. Untuk preprocess dataset Anda yang ditentukan sebagai file XYZ jalankan skrip preprocess_data.py . Contoh diberikan di sini:
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 Untuk melihat semua opsi dan sedikit deskripsi tentang mereka menjalankan python ./mace/scripts/preprocess_data.py --help . Skrip akan membuat sejumlah file HDF5 di folder processed_data yang dapat digunakan untuk pelatihan. Akan ada satu folder untuk pelatihan, satu untuk validasi dan yang terpisah untuk setiap config_type dalam set tes. Untuk melatih model, gunakan skrip run_train.py sebagai berikut:
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 Jika Anda ingin menggunakan MACE dengan bobot dan bias untuk mencatat eksperimen Anda cukup menginstal dengan
pip install ./mace[wandb] Dan tentukan argumen kata kunci yang diperlukan ( --wandb , --wandb_project , --wandb_entity , --wandb_name , --wandb_log_hypers )
Kami telah berkolaborasi dengan Proyek Bahan (MP) untuk melatih potensi roda universal yang mencakup 89 elemen pada kristal curah 1,6 M dalam dataset MPTRJ yang dipilih dari lintasan relaksasi MP. Model-model tersebut diulangi di GitHub di https://github.com/acuit/mace-mp. Jika Anda menggunakannya, silakan mengutip kertas kami yang juga berisi sejumlah besar aplikasi dan tolok ukur contoh.
Peringatan
Model MACE-MP dilatih pada energi DFT mentah MPTRJ dari output VASP, dan tidak secara langsung sebanding dengan energi DFT MP atau energi Chgnet, yang telah diterapkan Koreksi Kompatibilitas MP2020 untuk beberapa oksida logam transisi, fluorida (GGA/GGA+U Mixing untuk beberapa pencampuran logam transisi, GGA/GGA+U untuk pencampur u koreksi), dan 14 spesies anion (koreksi anion). Untuk detail lebih lanjut, silakan merujuk ke dokumentasi MP dan 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 ())Ada seri (kecil, sedang, besar) medan gaya organik yang dapat ditransfer. Ini dapat digunakan untuk simulasi molekul organik, kristal dan cairan molekul, atau sebagai titik awal untuk menyempurnakan pada dataset baru. Model dirilis di bawah lisensi ASL. Model-model tersebut diulangi di GitHub di https://github.com/acuit/mace-off. Jika Anda menggunakannya, silakan mengutip kertas kami yang juga berisi tolok ukur terperinci dan contoh aplikasi.
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 ()) Untuk memperbaiki salah satu model fondasi MACE-MP-0, Anda dapat menggunakan skrip mace_run_train dengan argumen tambahan --foundation_model=model_type . Misalnya untuk memperbaiki model kecil pada dataset baru, Anda dapat menggunakan:
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 Pilihan lain adalah "sedang" dan "besar", atau jalan menuju model fondasi. Jika Anda ingin finetune model lain, model akan dimuat dari jalur yang disediakan --foundation_model=$path_model , tetapi Anda perlu memberikan set penuh hyperparameters (irreps tersembunyi, r_max, dll.) Cocokkan model.
Proyek ini menggunakan pra-komit untuk menjalankan pemformatan kode dan berbaris pada komit. Kami juga menggunakan black , isort , pylint , dan mypy . Kami sarankan menyiapkan lingkungan pengembangan Anda dengan memasang paket dev ke lingkungan Python Anda:
pip install -e " .[dev] "
pre-commit install Baris kedua akan menginisialisasi pre-commit pra-untuk menjalankan pemeriksaan kode pada komit. Kami telah mengatur CI untuk memeriksa ini, tetapi kami sangat menyarankan Anda menjalankan perintah itu sebelum Anda melakukan (dan mendorong) untuk menghindari secara tidak sengaja melakukan kode buruk.
Kami senang menerima permintaan tarik di bawah lisensi MIT. Harap salin/tempel teks lisensi sebagai komentar ke permintaan tarik Anda.
Jika Anda menggunakan kode ini, silakan mengutip makalah kami:
@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 }
}Jika Anda memiliki pertanyaan, silakan hubungi kami di [email protected].
Untuk bug atau permintaan fitur, silakan gunakan masalah GitHub.
MACE diterbitkan dan didistribusikan di bawah lisensi MIT.
