Weitere Informationen zu dieser Arbeit finden Sie auf unserer Website.
Dieses Repository enthält den Code für den Artikel „From Molecules to Materials: Pre-training Large Generalizable Models for Atomic Property Prediction“.
Dieses Repository ist veraltet und wird nicht mehr aktiv gepflegt. Wir arbeiten derzeit daran, die Funktionalität dieses Repositorys in das offizielle Open Catalyst Project-Repository zu integrieren. Wenn Sie Fragen oder Bedenken zu diesem Repository haben, können Sie gerne ein Github-Problem erstellen oder uns per E-Mail kontaktieren. Bitte senden Sie eine E-Mail an Nima Shoghi und CC Brandon Wood.
In dieser Arbeit stellen wir Joint Multi-domain Pre-training (JMP) vor, eine überwachte Pre-Training-Strategie, die gleichzeitig auf mehreren Datensätzen aus verschiedenen chemischen Domänen trainiert und jeden Datensatz als einzigartige Pre-Training-Aufgabe innerhalb eines Multi-Task-Frameworks behandelt . Unser kombinierter Trainingsdatensatz besteht aus ca. 120 Millionen Systemen von OC20, OC22, ANI-1x und Transition-1x.
Die wichtigsten Beiträge dieser Arbeit sind:
Wir demonstrieren die leistungsstarke Generalisierungsfähigkeit von JMP, indem wir seine Feinabstimmungsleistung in einer vielfältigen Benchmark-Suite bewerten, die kleine Moleküle, große Moleküle und Materialien umfasst. JMP übertrifft durchweg das Training von Grund auf und erreicht oder erreicht den neuesten Stand bei 34 von 40 Feinabstimmungs-Benchmarks.
Wir zeigen, dass JMP eine effiziente Skalierung auf größere Modelle ermöglicht, die normalerweise überpassen würden, wenn sie von Grund auf auf kleinen Datensätzen trainiert würden. Das Vortraining fungiert als starker Regularisierer und ermöglicht es uns, ein Modell mit 235 Millionen Parametern zu trainieren, das bei mehreren Benchmarks mit wenigen Daten eine neue Leistung auf dem neuesten Stand der Technik setzt.
Wir führen eine detaillierte Analyse der Rechenanforderungen von JMP durch. Obwohl es im Vorfeld teuer ist, zeigen wir, dass sich die Vorschulungskosten von JMP durch eine über 12-mal schnellere Feinabstimmung im Vergleich zum Training von Grund auf amortisieren.
Wir glauben, dass JMP durch das Vortraining großer Modelle auf verschiedenen chemischen Daten einen wichtigen Schritt in Richtung des Ziels eines universellen ML-Potenzials für die Chemie darstellt. Das kontinuierliche Wachstum der verfügbaren Daten und der Rechenleistung wird die Fähigkeit von JMP, übertragbare atomare Darstellungen zu lernen, nur verbessern.
JMP weist eine durchschnittliche Verbesserung von 59 % im Vergleich zum Training von Grund auf auf und erreicht oder setzt bei 34 von 40 Aufgaben den neuesten Stand der Technik. Unsere Arbeit unterstreicht das Potenzial von Pre-Training-Strategien, die verschiedene Daten nutzen, um die Vorhersage von Eigenschaften über chemische Bereiche hinweg voranzutreiben, insbesondere für Aufgaben mit wenigen Daten.
Klonen Sie zunächst das Repository und navigieren Sie zum Stammverzeichnis:
git clone https://github.com/facebookresearch/JMP.git
cd JMP Richten Sie dann die Conda-Umgebung ein, wie in der Datei environment.yml bereitgestellt. Führen Sie dazu den folgenden Befehl aus (HINWEIS: Ersetzen Sie conda durch mamba , wenn Sie es installiert haben):
conda env create -f environment.yml -n jmpWenn der obige Befehl fehlschlägt, können Sie die Umgebung manuell selbst erstellen:
# 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 multiprocessAktivieren Sie dann die Umgebung:
conda activate jmpAbschließend installieren Sie das aktuelle Paket wie folgt:
pip install -e . Der Code ist nun einsatzbereit. Beispiele für die Ausführung des Codes finden Sie in den Konfigurationsdateien im Verzeichnis configs .
Die für das Vortraining und die Feinabstimmung verwendeten Daten sind aufgrund von Größenbeschränkungen nicht in diesem Repository enthalten. Nachfolgend finden Sie jedoch Anweisungen zum Herunterladen und Vorverarbeiten der Datensätze OC20, OC22, ANI-1x, Transition-1x, QM9, rMD17, MatBench, QMOF, SPICE und MD22.
Um den ANI-1x-Datensatz herunterzuladen und in ein Format zu konvertieren, das von unserer Codebasis verwendet werden kann, befolgen Sie bitte diese Schritte:
mkdir -p /path/to/datasets/ani1x
cd /path/to/datasets/ani1x.h5 -Format) von der offiziellen Quelle herunter: 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 -Dateien (setzen Sie --num_workers auf die Anzahl der verfügbaren CPU-Kerne): 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 Dateien in .lmdb -Dateien: 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.pklUm den Transition-1x-Datensatz herunterzuladen und in ein Format zu konvertieren, das von unserer Codebasis verwendet werden kann, führen Sie bitte die folgenden Schritte aus:
mkdir -p /path/to/datasets/transition1x
cd /path/to/datasets/transition1x.h5 -Format) von der offiziellen Quelle herunter: wget https://figshare.com/ndownloader/files/36035789 -O transition1x-release.h5.traj -Dateien (setzen Sie --num_workers auf die Anzahl der verfügbaren CPU-Kerne): 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 Dateien in .lmdb -Dateien: 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/ Die vorab trainierten Kontrollpunkte stehen unter den folgenden Links zum Download bereit:
Unsere Codebasis ist auf Python-Basis ausgelegt. Um ein Modell zu trainieren, erstellen Sie eine Konfigurationsdatei, die das Modell, den Datensatz und die Trainingsparameter angibt. Diese Konfigurationsobjekte werden mit Pydantic vollständig typgeprüft und validiert.
Sobald Sie ein Konfigurationsobjekt erstellt haben, können Sie die Klasse jmp.lightning.Runner verwenden, um die Trainingsschleife zu umschließen. Ein Beispiel für die Feinabstimmung eines Modells finden Sie unter configs/jmp_l_finetune.ipynb .
Der Großteil von JMP ist CC-BY-NC-lizenziert, wie in der LICENSE Datei zu finden ist. Teile des Projekts sind jedoch unter separaten Lizenzbedingungen verfügbar:
Wenn Sie diesen Code in Ihrer Forschung verwenden, zitieren Sie bitte das folgende Dokument:
@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}
}
