mace
v0.3.8
MACE fournit des potentiels interatomiques rapides et précis d'apprentissage automatique avec un message de message équivariant d'ordre supérieur.
Ce référentiel contient l'implémentation de référence MACE développée par Ilyes Batatia, Gregor Simm, David Kovacs et le groupe de Gabor Csanyi et des amis (voir contributeurs).
Également disponible:
Une documentation partielle est disponible sur: https://mace-docs.readthedocs.io
Assurez-vous d'installer Pytorch. Veuillez vous référer à l'installation officielle de Pytorch pour les instructions d'installation. Sélectionnez les options appropriées pour votre système.
C'est le moyen recommandé d'installer MACE.
pip install --upgrade pip
pip install mace-torchRemarque: Le package homonyme sur PYPI n'a rien à voir avec celui-ci.
git clone https://github.com/ACEsuit/mace.git
pip install ./mace Pour former un modèle MACE, vous pouvez utiliser le script mace_run_train , qui devrait être à l'endroit habituel que PIP place des binaires (ou vous pouvez explicitement exécuter 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 Pour donner un ensemble de validation spécifique, utilisez l'argument --valid_file . Pour définir une taille de lot plus grande pour évaluer l'ensemble de validation, spécifiez --valid_batch_size .
Pour contrôler la taille du modèle, vous devez changer --hidden_irreps . Pour la plupart des applications, la taille du modèle par défaut recommandée est --hidden_irreps='256x0e' (signifiant 256 messages invariants) ou --hidden_irreps='128x0e + 128x1o' . Si le modèle n'est pas suffisamment précis, vous pouvez inclure des fonctionnalités d'ordre supérieur, par exemple, 128x0e + 128x1o + 128x2e , ou augmenter le nombre de canaux à 256 . Il est également possible de spécifier le modèle à l'aide des touches --num_channels=128 et --max_L=1 .
Il est généralement préféré d'ajouter les atomes isolés à l'ensemble d'entraînement, plutôt que de lire dans leurs énergies à travers la ligne de commande comme dans l'exemple ci-dessus. Pour les étiqueter dans l'ensemble de formation, définissez config_type=IsolatedAtom dans leurs champs d'information. Si vous préférez ne pas utiliser ou ne connaissez pas les énergies des atomes isolés, vous pouvez utiliser l'option --E0s="average" qui estime les énergies atomiques en utilisant la régression des moindres carrés.
Si le mot-clé --swa est activé, le poids énergétique de la perte est augmenté pour les derniers 20% des époques d'entraînement (des époques --start_swa ). Ce paramètre aide généralement à réduire les erreurs d'énergie.
La précision peut être modifiée à l'aide du mot-clé --default_dtype , la valeur par défaut est float64 mais float32 donne une accélération significative (généralement un facteur de x2 dans la formation).
Les mots clés --batch_size et --max_num_epochs doivent être adaptés en fonction de la taille de l'ensemble de formation. La taille du lot doit être augmentée lorsque le nombre de données d'entraînement augmente et le nombre d'époches doit être diminué. Une heuristique pour les paramètres initiaux consiste à considérer le nombre de mise à jour de gradient constant à 200 000, qui peuvent être calculés comme
Le code peut gérer l'ensemble de formation avec des étiquettes hétérogènes, par exemple contenant les deux structures en vrac avec du stress et des molécules isolées. Dans cet exemple, pour faire en sorte que le code ignore le stress sur les molécules, ajoutez à votre configuration de molécules a config_stress_weight = 0.0 .
Pour utiliser l'Apple Silicon GPU Accélération, assurez-vous d'installer la dernière version Pytorch et spécifiez --device=mps .
Pour la formation multi-GPU, utilisez le drapeau --distributed . Cela utilisera le module DistributedDataparallel de Pytorch pour former le modèle sur plusieurs GPU. Combinez avec le chargement de données en ligne pour les grands ensembles de données (voir ci-dessous). Un exemple de script Slurm peut être trouvé dans mace/scripts/distributed_example.sbatch .
L'option pour analyser tous ou certains arguments à l'aide d'un YAML est disponible. Par exemple, pour former un modèle en utilisant les arguments ci-dessus, vous pouvez créer un fichier yaml your_configs.yaml avec le contenu suivant:
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
Et ajoutez à la ligne de commande --config="your_configs.yaml" . Tout argument spécifié dans la ligne de commande écrasera celui du fichier YAML.
Pour évaluer votre modèle MACE sur un fichier XYZ, exécutez le mace_eval_configs :
mace_eval_configs
--configs= " your_configs.xyz "
--model= " your_model.model "
--output= " ./your_output.xyz " Vous pouvez exécuter notre tutoriel Colab pour commencer rapidement avec Mace.
Nous avons également des tutoriels Colab plus détaillés sur:
Si vous disposez d'un grand ensemble de données qui pourrait ne pas s'adapter à la mémoire GPU, il est recommandé de prétraiter les données sur un CPU et d'utiliser la mise de données en ligne pour la formation du modèle. Pour prétraiter votre ensemble de données spécifié en tant que fichier XYZ, exécutez le script preprocess_data.py . Un exemple est donné ici:
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 Pour voir toutes les options et une petite description de leur exécution python ./mace/scripts/preprocess_data.py --help . Le script créera un certain nombre de fichiers HDF5 dans le dossier processed_data qui peuvent être utilisés pour la formation. Il y aura un dossier pour la formation, un pour la validation et un séparé pour chaque config_type dans l'ensemble de tests. Pour former le modèle, utilisez le script run_train.py comme suit:
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 Si vous souhaitez utiliser Mace avec des poids et des biais pour enregistrer vos expériences, il suffit d'installer avec
pip install ./mace[wandb] Et spécifiez les arguments de mots clés nécessaires ( --wandb , --wandb_project , --wandb_entity , --wandb_name , --wandb_log_hypers )
Nous avons collaboré avec le projet de matériaux (MP) pour former un potentiel de masse universel couvrant 89 éléments sur des cristaux en vrac de 1,6 m dans l'ensemble de données MPTRJ sélectionné dans les trajectoires de relaxation MP. Les modèles sont liés sur GitHub sur https://github.com/acesiuit/mace-mp. Si vous les utilisez, veuillez citer notre papier qui contient également une large gamme d'exemples d'applications et de repères.
Prudence
Les modèles MACE-MP sont formés sur les énergies DFT RAW MPTRJ à partir des sorties VASP, et ne sont pas directement comparables aux énergies DFT du MP ou aux énergies de Chgnet, qui ont été appliquées MP20202 Corrections) et 14 espèces anions (corrections d'anions). Pour plus de détails, veuillez vous référer à la documentation MP et à MP2020CAPATibilité.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 ())Il y a une série (petits, moyens, grands) champs de force organique transférables. Ceux-ci peuvent être utilisés pour la simulation de molécules organiques, de cristaux et de liquides moléculaires, ou comme point de départ pour le réglage fin sur un nouvel ensemble de données. Les modèles sont publiés sous la licence ASL. Les modèles sont liés sur GitHub sur https://github.com/acesity/mace-off. Si vous les utilisez, veuillez citer notre article qui contient également des repères détaillés et des exemples d'applications.
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 ()) Pour Finetune un du modèle de fondation MACE-MP-0, vous pouvez utiliser le script mace_run_train avec l'argument supplémentaire --foundation_model=model_type . Par exemple, pour Finetune le petit modèle sur un nouvel ensemble de données, vous pouvez utiliser:
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 Les autres options sont "moyennes" et "grandes", ou le chemin vers un modèle de fondation. Si vous souhaitez FineTune un autre modèle, le modèle sera chargé à partir du chemin fourni --foundation_model=$path_model , mais vous devrez fournir l'ensemble complet d'hyperparamètres (Irreps cachés, R_Max, etc.) correspondant au modèle.
Ce projet utilise un pré-engagement pour exécuter le formatage de code et la libellur sur le commit. Nous utilisons également black , isort , pylint et mypy . Nous vous recommandons de configurer votre environnement de développement en installant les packages dev dans votre environnement Python:
pip install -e " .[dev] "
pre-commit install La deuxième ligne initialisera pre-commit à des vérifications de code d'exécution automaticales sur le commit. Nous avons configuré CI pour vérifier cela, mais nous vous recommandons fortement d'exécuter ces commandes avant de vous engager (et de pousser) pour éviter de commettre accidentellement un mauvais code.
Nous sommes heureux d'accepter les demandes de traction sous une licence du MIT. Veuillez copier / coller le texte de licence en tant que commentaire dans votre demande de traction.
Si vous utilisez ce code, veuillez citer nos papiers:
@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 }
}Si vous avez des questions, veuillez nous contacter à [email protected].
Pour les bogues ou les demandes de fonctionnalités, veuillez utiliser les problèmes GitHub.
MACE est publié et distribué sous la licence du MIT.
