Veuillez consulter notre site Web pour plus d’informations sur ce travail.
Ce référentiel contient le code de l'article « Des molécules aux matériaux : pré-formation de grands modèles généralisables pour la prévision des propriétés atomiques ».
Ce référentiel est obsolète et n'est plus activement maintenu. Nous travaillons actuellement sur l'intégration des fonctionnalités de ce référentiel dans le référentiel officiel de l'Open Catalyst Project. Si vous avez des questions ou des préoccupations concernant ce référentiel, n'hésitez pas à créer un problème github ou à nous contacter par e-mail. Veuillez envoyer un e-mail à Nima Shoghi et CC Brandon Wood.
Dans ce travail, nous introduisons le Joint Multi-domain Pre-training (JMP), une stratégie de pré-entraînement supervisé qui s'entraîne simultanément sur plusieurs ensembles de données provenant de différents domaines chimiques, traitant chaque ensemble de données comme une tâche de pré-entraînement unique dans un cadre multi-tâches. . Notre ensemble de données de formation combinées comprend environ 120 millions de systèmes d'OC20, OC22, ANI-1x et Transition-1x.
Les principales contributions de ce travail sont :
Nous démontrons la puissante capacité de généralisation de JMP en évaluant ses performances de réglage précis sur une suite diversifiée de références couvrant de petites molécules, de grandes molécules et des matériaux. JMP surpasse systématiquement la formation à partir de zéro et établit ou correspond à l'état de l'art sur 34 des 40 tests de réglage fin.
Nous montrons que JMP permet une mise à l'échelle efficace vers des modèles plus grands qui seraient normalement surajustés s'ils étaient entraînés à partir de zéro sur de petits ensembles de données. La pré-formation agit comme un puissant régulateur, nous permettant de former un modèle avec 235 millions de paramètres qui établit de nouvelles performances de pointe sur plusieurs benchmarks à faible quantité de données.
Nous effectuons une analyse détaillée des exigences informatiques de JMP. Bien que coûteux au départ, nous montrons que le coût de pré-formation de JMP est récupéré en permettant un réglage précis 12 fois plus rapide qu'une formation à partir de zéro.
En pré-entraînant de grands modèles sur diverses données chimiques, nous pensons que JMP représente une étape importante vers l’objectif d’un potentiel universel de ML pour la chimie. La croissance continue des données disponibles et de la puissance de calcul ne fera qu'améliorer la capacité de JMP à apprendre des représentations atomiques transférables.
JMP démontre une amélioration moyenne de 59 % par rapport à une formation à partir de zéro et correspond ou définit l'état de l'art sur 34 tâches sur 40. Nos travaux mettent en évidence le potentiel des stratégies de pré-formation qui utilisent diverses données pour faire progresser la prédiction des propriétés dans tous les domaines chimiques, en particulier pour les tâches nécessitant peu de données.
Tout d’abord, clonez le référentiel et accédez au répertoire racine :
git clone https://github.com/facebookresearch/JMP.git
cd JMP Ensuite, configurez l'environnement conda, comme indiqué dans le fichier environment.yml . Pour ce faire, exécutez la commande suivante (REMARQUE : remplacez conda par mamba si vous l'avez installé) :
conda env create -f environment.yml -n jmpSi la commande ci-dessus échoue, vous pouvez créer vous-même l'environnement manuellement :
# 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 multiprocessEnsuite, activez l'environnement :
conda activate jmpEnfin, installez le package actuel comme suit :
pip install -e . Le code est maintenant prêt à être utilisé. Consultez les fichiers de configuration dans le répertoire configs pour obtenir des exemples sur la façon d'exécuter le code.
Les données utilisées pour la pré-formation et le réglage fin ne sont pas incluses dans ce référentiel en raison de contraintes de taille. Cependant, les instructions de téléchargement et de prétraitement des ensembles de données OC20, OC22, ANI-1x, Transition-1x, QM9, rMD17, MatBench, QMOF, SPICE et MD22 sont fournies ci-dessous.
Pour télécharger l'ensemble de données ANI-1x et le convertir dans un format pouvant être utilisé par notre base de code, veuillez suivre ces étapes :
mkdir -p /path/to/datasets/ani1x
cd /path/to/datasets/ani1x.h5 ) depuis la source officielle : 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 (définissez --num_workers sur le nombre de cœurs de processeur dont vous disposez) : 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 en fichiers .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.pklPour télécharger l'ensemble de données Transition-1x et le convertir dans un format pouvant être utilisé par notre base de code, veuillez suivre ces étapes :
mkdir -p /path/to/datasets/transition1x
cd /path/to/datasets/transition1x.h5 ) depuis la source officielle : wget https://figshare.com/ndownloader/files/36035789 -O transition1x-release.h5.traj (définissez --num_workers sur le nombre de cœurs de processeur dont vous disposez) : 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 en fichiers .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/ Les points de contrôle pré-entraînés sont disponibles en téléchargement à partir des liens suivants :
Notre base de code est conçue pour être basée sur Python. Pour entraîner un modèle, vous créez un fichier de configuration qui spécifie le modèle, l'ensemble de données et les paramètres d'entraînement. Ces objets de configuration sont entièrement vérifiés et validés à l'aide de Pydantic.
Une fois que vous avez créé un objet de configuration, vous pouvez utiliser la classe jmp.lightning.Runner pour encapsuler la boucle de formation. Voir configs/jmp_l_finetune.ipynb pour un exemple de la façon d'affiner un modèle.
La majorité de JMP est sous licence CC-BY-NC, comme indiqué dans le fichier LICENSE . Cependant, certaines parties du projet sont disponibles sous des conditions de licence distinctes :
Si vous utilisez ce code dans votre recherche, veuillez citer l'article suivant :
@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}
}
