mace
v0.3.8
MACE proporciona potenciales interatómicos de aprendizaje automático rápido y preciso con un mensaje equivalente de mensajes de orden superior.
Este repositorio contiene la implementación de referencia de MACE desarrollada por Ilyes Batatia, Gregor Simm, David Kovacs y el grupo de Gabor Csanyi y amigos (ver contribuyentes).
También disponible:
Una documentación parcial está disponible en: https://mace-docs.readthedocs.io
Asegúrese de instalar Pytorch. Consulte la instalación oficial de Pytorch para las instrucciones de instalación. Seleccione las opciones apropiadas para su sistema.
Esta es la forma recomendada de instalar MACE.
pip install --upgrade pip
pip install mace-torchNota: El paquete homónimo en Pypi no tiene nada que ver con este.
git clone https://github.com/ACEsuit/mace.git
pip install ./mace Para entrenar un modelo MACE, puede usar el script mace_run_train , que debería estar en el lugar habitual que Pip coloca binarios (o puede ejecutar explícitamente 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 Para dar un conjunto de validación específico, use el argumento --valid_file . Para establecer un tamaño de lote más grande para evaluar el conjunto de validación, especifique --valid_batch_size .
Para controlar el tamaño del modelo, debe cambiar --hidden_irreps . Para la mayoría de las aplicaciones, el tamaño del modelo predeterminado recomendado es --hidden_irreps='256x0e' (que significa 256 mensajes invariantes) o --hidden_irreps='128x0e + 128x1o' . Si el modelo no es lo suficientemente preciso, puede incluir características de orden superior, por ejemplo, 128x0e + 128x1o + 128x2e , o aumentar el número de canales a 256 . También es posible especificar el modelo usando las teclas --num_channels=128 y --max_L=1 .
Por lo general, se prefiere agregar los átomos aislados al conjunto de entrenamiento, en lugar de leer en sus energías a través de la línea de comando como en el ejemplo anterior. Para etiquetarlos en el conjunto de entrenamiento, establezca config_type=IsolatedAtom en sus campos de información. Si prefiere no usar o no conocer las energías de los átomos aislados, puede usar la opción --E0s="average" que estima las energías atómicas utilizando la regresión de mínimos cuadrados.
Si la palabra clave --swa está habilitada, el peso energético de la pérdida aumenta para el último ~ 20% de las épocas de entrenamiento (de las épocas --start_swa ). Esta configuración generalmente ayuda a reducir los errores de energía.
La precisión se puede cambiar utilizando la palabra clave --default_dtype , el valor predeterminado es float64 pero float32 ofrece una aceleración significativa (generalmente un factor de x2 en el entrenamiento).
Las palabras clave --batch_size y --max_num_epochs deben adaptarse en función del tamaño del conjunto de entrenamiento. El tamaño del lote debe aumentarse cuando aumenta el número de datos de entrenamiento, y el número de épocas debe disminuir. Una heurística para la configuración inicial es considerar el número de actualización de gradiente constante a 200 000, que se puede calcular como
El código puede manejar el conjunto de entrenamiento con etiquetas heterogéneas, por ejemplo, que contiene ambas estructuras a granel con estrés y moléculas aisladas. En este ejemplo, para hacer que el código ignore el estrés en las moléculas, agregue la configuración de las moléculas una config_stress_weight = 0.0 .
Para usar la aceleración de GPU de Silicon Apple, asegúrese de instalar la última versión de Pytorch y especificar --device=mps .
Para capacitación multi-GPU, use la bandera --distributed . Esto utilizará el módulo DistributedDataparallel DistributedDataparallel de Pytorch para entrenar el modelo en múltiples GPU. Combine con la carga de datos en línea para conjuntos de datos grandes (ver más abajo). Se puede encontrar un script de Slurm de ejemplo en mace/scripts/distributed_example.sbatch .
La opción para analizar todos o algunos argumentos utilizando un YAML está disponible. Por ejemplo, para entrenar un modelo que usa los argumentos anteriores, puede crear un archivo YAML your_configs.yaml con el siguiente contenido:
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
Y agregar a la línea de comando --config="your_configs.yaml" . Cualquier argumento especificado en la línea de comandos sobrescribirá el del archivo YAML.
Para evaluar su modelo MACE en un archivo XYZ, ejecute mace_eval_configs :
mace_eval_configs
--configs= " your_configs.xyz "
--model= " your_model.model "
--output= " ./your_output.xyz " Puede ejecutar nuestro tutorial de Colab para comenzar rápidamente con Mace.
También tenemos tutoriales de Colab más detallados en:
Si tiene un conjunto de datos grande que podría no encajar en la memoria de la GPU, se recomienda preprocesar los datos en una CPU y usar DataLoading en línea para capacitar al modelo. Para preprocesar su conjunto de datos especificado como un archivo XYZ, ejecute el script preprocess_data.py . Aquí se da un ejemplo:
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 Para ver todas las opciones y una pequeña descripción de ellas, ejecutan python ./mace/scripts/preprocess_data.py --help . El script creará una serie de archivos HDF5 en la carpeta processed_data que se puede utilizar para la capacitación. Habrá una carpeta para el entrenamiento, una para la validación y una separada para cada config_type en el conjunto de pruebas. Para entrenar el modelo, use el script run_train.py de la siguiente manera:
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 desea usar Mace con pesas y sesgos para registrar sus experimentos, simplemente instale con
pip install ./mace[wandb] Y especifique los argumentos de palabras clave necesarias ( --wandb , --wandb_project , --wandb_entity , --wandb_name , --wandb_log_hypers )
Hemos colaborado con el Proyecto de Materiales (MP) para capacitar un potencial de maza universal que cubre 89 elementos en cristales a granel de 1,6 m en el conjunto de datos MPTRJ seleccionado de las trayectorias de relajación de MP. Los modelos se relajan en GitHub en https://github.com/acessuit/mace-mp. Si los usa, cite nuestro documento que también contiene una amplia gama de aplicaciones y puntos de referencia de ejemplo.
Precaución
Los modelos MACE-MP están entrenados en energías DFT sin procesar MPTRJ de salidas VASP, y no son directamente comparables a las energías DFT del MP o las energías de Chgnet, que se han aplicado correcciones de competencia MP2020 para algunos óxidos de metales de transición, fluoruros (GGA/GGA+U Mezcla correcciones), y 14 especies de aniones (correcciones de aniones). Para obtener más detalles, consulte la documentación de MP y 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 ())Hay una serie (pequeño, mediano, grande) campos de fuerza orgánica transferible. Estos se pueden usar para la simulación de moléculas orgánicas, cristales y líquidos moleculares, o como punto de partida para ajustar un nuevo conjunto de datos. Los modelos se lanzan bajo la licencia ASL. Los modelos se relajan en GitHub en https://github.com/acessuit/mace-off. Si los usa, cite nuestro documento que también contiene puntos de referencia detallados y aplicaciones de ejemplo.
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 ()) Para Finetune uno del modelo MACE-MP-0 Foundation, puede usar el script mace_run_train con el argumento adicional --foundation_model=model_type . Por ejemplo, para Finetune el modelo pequeño en un nuevo conjunto de datos, puede usar:
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 Otras opciones son "medianas" y "grandes", o el camino hacia un modelo de base. Si desea Finetune otro modelo, el modelo se cargará desde la ruta proporcionada --foundation_model=$path_model , pero deberá proporcionar el conjunto completo de hiperparámetros (Irreps ocultos, R_max, etc.) que coincida con el modelo.
Este proyecto utiliza pre-Commit para ejecutar el formato de código y la pelusa en la confirmación. También usamos black , isort , pylint y mypy . Recomendamos configurar su entorno de desarrollo instalando los paquetes dev en su entorno Python:
pip install -e " .[dev] "
pre-commit install La segunda línea inicializará pre-commit para ejecutar las verificaciones de código de ejecución automática en la confirmación. Tenemos CI configurado para verificar esto, pero le recomendamos que ejecute esos comandos antes de comprometerse (y presionar) para evitar cometer un código incorrecto accidentalmente.
Nos complace aceptar solicitudes de extracción bajo una licencia MIT. Copie/pegue el texto de la licencia como un comentario en su solicitud de extracción.
Si usa este código, cite nuestros documentos:
@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 tiene alguna pregunta, contáctenos en [email protected].
Para errores o solicitudes de funciones, utilice problemas de GitHub.
Mace se publica y distribuye bajo la licencia MIT.
