mace
v0.3.8
MACE обеспечивает быстрые и точные межатомные потенциалы машинного обучения с прохождением эквивалентного сообщения более высокого порядка.
Этот репозиторий содержит справочную реализацию 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 . В нем будет использоваться DistributedDataparallel модуля Pytorch для обучения модели на нескольких графических процессорах. В сочетании с загрузкой данных в Интернете для больших наборов данных (см. Ниже). Пример сценария 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.
Чтобы оценить свою модель MACE в файле XYZ, запустите mace_eval_configs :
mace_eval_configs
--configs= " your_configs.xyz "
--model= " your_model.model "
--output= " ./your_output.xyz " Вы можете запустить наш учебник Colab, чтобы быстро начать с MACE.
У нас также есть более подробные учебники COLAB на:
Если у вас есть большой набор данных, который может не вписаться в память графического процессора, рекомендуется предварительно обработать данные на процессоре и использовать онлайн-загрузку данных для обучения модели. Чтобы предварительно обработать ваш набор данных, указанный в виде файла 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 . Сценарий создаст несколько файлов HDF5 в папке processed_data , которые можно использовать для обучения. Там будет одна папка для обучения, одна для проверки и отдельная для каждого 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 Если вы хотите использовать MACE с весами и смещениями для регистрации ваших экспериментов.
pip install ./mace[wandb] И укажите необходимые аргументы ключевых слов ( --wandb , --wandb_project , --wandb_entity , --wandb_name , --wandb_log_hypers )
Мы сотрудничали с проектом «Материалы» (MP) для обучения универсального потенциала MACE, охватывающего 89 элементов на 1,6 м массовых кристаллов в наборе данных MPTRJ, выбранном из траекторий релаксации MP. Модели переселены на GitHub по адресу https://github.com/aceuit/mace-mp. Если вы используете их, пожалуйста, процитируйте нашу статью, которая также содержит большой диапазон примеров приложений и тестов.
Осторожность
Модели MACE-MP обучаются на MPTRJ RAW Energies DFT из выходов VASP и не имеют напрямую сопоставимы с энергиями DFT MP или энергиями CHGNET, которые были применены коррекции MP202020202020202020202 исправления) и 14 видов анионов (анионные коррекции). Для получения более подробной информации, пожалуйста, обратитесь к документации MP и MP20202020202020202020202020202020202020202020202020202020202020202020202020
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. Модели переселены на GitHub по адресу https://github.com/aceuit/mace-ft. Если вы используете их, пожалуйста, процитируйте нашу статью, которая также содержит подробные показатели и примеры приложений.
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, вы можете использовать скрипт mace_run_train с дополнительным аргументом --foundation_model=model_type . Например, чтобы определить небольшую модель на новом наборе данных, вы можете использовать:
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 Другими вариантами являются «средние» и «большой» или путь к модели фундамента. Если вы хотите получить еще одну модель, модель будет загружена из предоставленного пути --foundation_model=$path_model , но вам необходимо будет обеспечить полный набор гиперпараметра (скрытые неврепы, r_max и т. Д.), Сопоставление модели.
Этот проект использует предварительную компанию для выполнения форматирования и снятия кода на коммите. Мы также используем black , isort , pylint и mypy . Мы рекомендуем настроить вашу среду разработки, установив пакеты dev в среду Python:
pip install -e " .[dev] "
pre-commit install Вторая строка будет инициализировать 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.
