Пожалуйста, посетите наш сайт для получения дополнительной информации об этой работе.
Этот репозиторий содержит код статьи «От молекул к материалам: предварительное обучение больших обобщаемых моделей для прогнозирования свойств атомов».
Этот репозиторий устарел и больше не поддерживается. В настоящее время мы работаем над интеграцией функциональности этого репозитория в официальный репозиторий Open Catalyst Project. Если у вас есть какие-либо вопросы или сомнения относительно этого репозитория, пожалуйста, создайте проблему на GitHub или свяжитесь с нами по электронной почте. Пожалуйста, отправьте электронное письмо Ниме Шоги и CC Брэндону Вуду.
В этой работе мы представляем совместное многодоменное предварительное обучение (JMP), контролируемую стратегию предварительного обучения, которая одновременно обучает нескольким наборам данных из разных химических областей, рассматривая каждый набор данных как уникальную задачу предварительного обучения в многозадачной структуре. . Наш объединенный набор обучающих данных состоит из примерно 120 миллионов систем OC20, OC22, ANI-1x и Transition-1x.
Ключевыми вкладами этой работы являются:
Мы демонстрируем мощную способность JMP к обобщению, оценивая его производительность в рамках разнообразного набора тестов, охватывающего малые молекулы, большие молекулы и материалы. JMP постоянно превосходит обучение с нуля и устанавливает или соответствует современному уровню по 34 из 40 тестов тонкой настройки.
Мы показываем, что JMP обеспечивает эффективное масштабирование до более крупных моделей, которые обычно переобучаются при обучении с нуля на небольших наборах данных. Предварительное обучение действует как сильный регуляризатор, позволяя нам обучать модель с 235 миллионами параметров, которые обеспечивают новую современную производительность на нескольких тестах с низким объемом данных.
Мы проводим подробный анализ вычислительных требований JMP. Несмотря на то, что первоначальные затраты стоят дорого, мы показываем, что затраты на предварительное обучение JMP окупаются за счет обеспечения более чем в 12 раз более быстрой тонкой настройки по сравнению с обучением с нуля.
Мы считаем, что JMP, предварительно обучая большие модели на разнообразных химических данных, представляет собой важный шаг на пути к созданию универсального потенциала машинного обучения для химии. Продолжающийся рост доступных данных и вычислительной мощности только улучшит способность JMP изучать передаваемые атомарные представления.
JMP демонстрирует среднее улучшение на 59% по сравнению с обучением с нуля и соответствует или устанавливает самые современные решения в 34 из 40 задач. Наша работа подчеркивает потенциал стратегий предварительного обучения, которые используют разнообразные данные для улучшения прогнозирования свойств в различных химических областях, особенно для задач с низким объемом данных.
Сначала клонируйте репозиторий и перейдите в корневой каталог:
git clone https://github.com/facebookresearch/JMP.git
cd JMP Затем настройте среду conda, как указано в файле environment.yml . Для этого выполните следующую команду (ПРИМЕЧАНИЕ: замените conda на mamba если она у вас установлена):
conda env create -f environment.yml -n jmpЕсли приведенная выше команда не удалась, вы можете создать среду вручную самостоятельно:
# Create the environment
conda create -n jmp python=3.11
conda activate jmp
# Install PyTorch
conda install -y -c pytorch -c nvidia pytorch torchvision torchaudio pytorch-cuda=12.1
# Install PyG, PyTorch Scatter, PyTorch Sparse.
conda install -c pyg pyg pytorch-sparse pytorch-cluster
# Install other conda dependencies
conda install -y
-c conda-forge
numpy matplotlib seaborn sympy pandas numba scikit-learn plotly nbformat ipykernel ipywidgets tqdm pyyaml networkx
pytorch-lightning torchmetrics lightning
einops wandb
cloudpickle
" pydantic>2 "
frozendict wrapt varname typing-extensions lovely-tensors lovely-numpy requests pytest nbval
# Install pip dependencies
pip install lmdb
# Install dependencies for materials datasets
pip install ase
# Install dependencies for large molecule datasets
conda install h5py
# Install MatBench dependencies
pip install matbench
# Install dependencies for PDBBind
pip install biopython rdkit
# Install dependencies for pre-processing ANI1x/Transition1x
pip install multiprocessЗатем активируйте среду:
conda activate jmpНаконец, установите текущий пакет следующим образом:
pip install -e . Теперь код готов к использованию. См. файлы конфигурации в каталоге configs , где приведены примеры запуска кода.
Данные, используемые для предварительного обучения и точной настройки, не включены в этот репозиторий из-за ограничений по размеру. Однако инструкции по загрузке и предварительной обработке наборов данных OC20, OC22, ANI-1x, Transition-1x, QM9, rMD17, MatBench, QMOF, SPICE и MD22 приведены ниже.
Чтобы загрузить набор данных ANI-1x и преобразовать его в формат, который может использоваться нашей кодовой базой, выполните следующие действия:
mkdir -p /path/to/datasets/ani1x
cd /path/to/datasets/ani1x.h5 ) из официального источника: wget https://springernature.figshare.com/ndownloader/files/18112775 -O ani1x-release.h5python -m jmp.datasets.scripts.ani1x_preprocess.ani1x_splits --input_file ani1x-release.h5 --train_keys_output train_keys.pkl --val_keys_output val_keys.pkl --test_keys_output test_keys.pkl.traj (задайте --num_workers количество доступных ядер ЦП): mkdir -p traj
mkdir -p traj/train traj/val traj/test
python -m jmp.datasets.scripts.ani1x_preprocess.ani1x_write_traj --ani1x_h5 ani1x-release.h5 --split_keys train_keys.pkl --split train --traj_dir traj/train --num_workers 32
python -m jmp.datasets.scripts.ani1x_preprocess.ani1x_write_traj --ani1x_h5 ani1x-release.h5 --split_keys val_keys.pkl --split val --traj_dir traj/val --num_workers 32
python -m jmp.datasets.scripts.ani1x_preprocess.ani1x_write_traj --ani1x_h5 ani1x-release.h5 --split_keys test_keys.pkl --split test --traj_dir traj/test --num_workers 32.traj в файлы .lmdb : mkdir -p lmdb
mkdir -p lmdb/train lmdb/val lmdb/test
python -m jmp.datasets.scripts.ani1x_preprocess.ani1x_write_lmdbs --data_path traj/train --out_path lmdb/train --split train --num_workers 32
python -m jmp.datasets.scripts.ani1x_preprocess.ani1x_write_lmdbs --data_path traj/val --out_path lmdb/val --split val --num_workers 32
python -m jmp.datasets.scripts.ani1x_preprocess.ani1x_write_lmdbs --data_path traj/test --out_path lmdb/test --split test --num_workers 32python -m jmp.datasets.scripts.ani1x_preprocess.ani1x_linear_ref linref --src lmdb/train --out_path linref.npzpython -m jmp.datasets.scripts.ani1x_preprocess.ani1x_linear_ref compute_mean_std --src lmdb/train --linref_path linref.npz --out_path mean_std.pklЧтобы загрузить набор данных Transition-1x и преобразовать его в формат, который может использоваться нашей кодовой базой, выполните следующие действия:
mkdir -p /path/to/datasets/transition1x
cd /path/to/datasets/transition1x.h5 ) из официального источника: wget https://figshare.com/ndownloader/files/36035789 -O transition1x-release.h5.traj (задайте --num_workers количество доступных ядер ЦП): mkdir -p traj
mkdir -p traj/train traj/val traj/test
python -m jmp.datasets.scripts.transition1x_preprocess.trans1x_write_traj --transition1x_h5 transition1x-release.h5 --split train --traj_dir traj/train --num_workers 32
python -m jmp.datasets.scripts.transition1x_preprocess.trans1x_write_traj --transition1x_h5 transition1x-release.h5 --split val --traj_dir traj/val --num_workers 32
python -m jmp.datasets.scripts.transition1x_preprocess.trans1x_write_traj --transition1x_h5 transition1x-release.h5 --split test --traj_dir traj/test --num_workers 32.traj в файлы .lmdb : mkdir -p lmdb
mkdir -p lmdb/train lmdb/val lmdb/test
python -m jmp.datasets.scripts.transition1x_preprocess.trans1x_write_lmdbs --data_path traj/train --out_path lmdb/train --split train --num_workers 32
python -m jmp.datasets.scripts.transition1x_preprocess.trans1x_write_lmdbs --data_path traj/val --out_path lmdb/val --split val --num_workers 32
python -m jmp.datasets.scripts.transition1x_preprocess.trans1x_write_lmdbs --data_path traj/test --out_path lmdb/test --split test --num_workers 32python -m jmp.datasets.scripts.transition1x_preprocess.trans1x_linear_ref linref --src lmdb/train --out_path linref.npzpython -m jmp.datasets.scripts.transition1x_preprocess.trans1x_linear_ref compute_mean_std --src lmdb/train --linref_path linref.npz --out_path mean_std.pklpython -m jmp.datasets.finetune.rmd17 download --destination /path/to/datasets/rmd17/python -m jmp.datasets.finetune.qm9 download --destination /path/to/datasets/qm9/python -m jmp.datasets.finetune.md22 download --destination /path/to/datasets/md22/python -m jmp.datasets.finetune.spice download --destination /path/to/datasets/spice/python -m jmp.datasets.finetune.mat_bench download --destination /path/to/datasets/matbench/python -m jmp.datasets.finetune.qmof download --destination /path/to/datasets/qmof/ Предварительно обученные контрольные точки доступны для загрузки по следующим ссылкам:
Наша кодовая база спроектирована на основе Python. Чтобы обучить модель, вы создаете файл конфигурации, в котором указываются модель, набор данных и параметры обучения. Эти объекты конфигурации полностью проверяются и проверяются с помощью Pydantic.
После создания объекта конфигурации вы можете использовать класс jmp.lightning.Runner для завершения цикла обучения. См. configs/jmp_l_finetune.ipynb для примера точной настройки модели.
Большая часть JMP имеет лицензию CC-BY-NC, как указано в файле LICENSE . Однако части проекта доступны по отдельным условиям лицензии:
Если вы используете этот код в своих исследованиях, цитируйте следующую статью:
@article{shoghi2023molecules,
title={From molecules to materials: Pre-training large generalizable models for atomic property prediction},
author={Shoghi, Nima and Kolluru, Adeesh and Kitchin, John R and Ulissi, Zachary W and Zitnick, C Lawrence and Wood, Brandon M},
journal={arXiv preprint arXiv:2310.16802},
year={2023}
}
