Consulte nuestro sitio web para obtener más información sobre este trabajo.
Este repositorio contiene el código del artículo "De moléculas a materiales: entrenamiento previo de grandes modelos generalizables para la predicción de propiedades atómicas".
Este repositorio está en desuso y ya no se mantiene activamente. Actualmente estamos trabajando para integrar la funcionalidad de este repositorio en el repositorio oficial de Open Catalyst Project. Si tiene alguna pregunta o inquietud con respecto a este repositorio, no dude en crear un problema de github o comunicarse con nosotros por correo electrónico. Envíe un correo electrónico a Nima Shoghi y CC Brandon Wood.
En este trabajo, presentamos el entrenamiento previo conjunto multidominio (JMP), una estrategia de entrenamiento previo supervisado que entrena simultáneamente en múltiples conjuntos de datos de diferentes dominios químicos, tratando cada conjunto de datos como una tarea de entrenamiento previo única dentro de un marco de tareas múltiples. . Nuestro conjunto de datos de entrenamiento combinado consta de ~120 millones de sistemas de OC20, OC22, ANI-1x y Transition-1x.
Las contribuciones clave de este trabajo son:
Demostramos la poderosa capacidad de generalización de JMP al evaluar su rendimiento de ajuste en un conjunto diverso de puntos de referencia que abarca moléculas pequeñas, moléculas grandes y materiales. JMP supera consistentemente el entrenamiento desde cero y establece o iguala lo último en 34 de 40 puntos de referencia de ajuste.
Mostramos que JMP permite un escalamiento eficiente a modelos más grandes que normalmente se sobreajustarían si se entrenaran desde cero en conjuntos de datos pequeños. El entrenamiento previo actúa como un potente regularizador, lo que nos permite entrenar un modelo con parámetros de 235 M que establece un rendimiento de última generación en múltiples puntos de referencia con pocos datos.
Realizamos un análisis detallado de los requisitos computacionales de JMP. Si bien es costoso desde el principio, demostramos que el costo previo a la capacitación de JMP se recupera al permitir un ajuste fino 12 veces más rápido en comparación con la capacitación desde cero.
Al entrenar previamente modelos grandes con diversos datos químicos, creemos que JMP representa un paso importante hacia el objetivo de un potencial de aprendizaje automático universal para la química. El crecimiento continuo de los datos disponibles y la potencia informática no hará más que mejorar la capacidad de JMP para aprender representaciones atómicas transferibles.
JMP demuestra una mejora promedio del 59 % con respecto al entrenamiento desde cero, e iguala o establece lo último en 34 de 40 tareas. Nuestro trabajo destaca el potencial de las estrategias de preentrenamiento que utilizan datos diversos para avanzar en la predicción de propiedades en dominios químicos, especialmente para tareas con pocos datos.
Primero, clone el repositorio y navegue hasta el directorio raíz:
git clone https://github.com/facebookresearch/JMP.git
cd JMP Luego, configure el entorno conda, como se proporciona en el archivo environment.yml . Para hacerlo, ejecute el siguiente comando (NOTA: reemplace conda con mamba si lo tiene instalado):
conda env create -f environment.yml -n jmpSi el comando anterior falla, puedes crear el entorno manualmente tú mismo:
# 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 multiprocessLuego, activa el entorno:
conda activate jmpFinalmente, instale el paquete actual de la siguiente manera:
pip install -e . El código ya está listo para ser utilizado. Consulte los archivos de configuración en el directorio configs para obtener ejemplos sobre cómo ejecutar el código.
Los datos utilizados para la capacitación previa y el ajuste no se incluyen en este repositorio debido a limitaciones de tamaño. Sin embargo, a continuación se proporcionan instrucciones para descargar y preprocesar los conjuntos de datos OC20, OC22, ANI-1x, Transition-1x, QM9, rMD17, MatBench, QMOF, SPICE y MD22.
Para descargar el conjunto de datos ANI-1x y convertirlo a un formato que pueda ser utilizado por nuestro código base, siga estos pasos:
mkdir -p /path/to/datasets/ani1x
cd /path/to/datasets/ani1x.h5 ) de la fuente oficial: 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 (establezca --num_workers en la cantidad de núcleos de CPU que tenga disponibles): 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 a archivos .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.pklPara descargar el conjunto de datos Transition-1x y convertirlo a un formato que pueda ser utilizado por nuestro código base, siga estos pasos:
mkdir -p /path/to/datasets/transition1x
cd /path/to/datasets/transition1x.h5 ) de la fuente oficial: wget https://figshare.com/ndownloader/files/36035789 -O transition1x-release.h5.traj (establezca --num_workers en la cantidad de núcleos de CPU que tenga disponibles): 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 a archivos .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/ Los puntos de control previamente entrenados están disponibles para descargar desde los siguientes enlaces:
Nuestra base de código está diseñada para estar basada en Python. Para entrenar un modelo, crea un archivo de configuración que especifica el modelo, el conjunto de datos y los parámetros de entrenamiento. Estos objetos de configuración se verifican y validan completamente mediante Pydantic.
Una vez que haya creado un objeto de configuración, puede usar la clase jmp.lightning.Runner para ajustar el ciclo de entrenamiento. Consulte configs/jmp_l_finetune.ipynb para ver un ejemplo de cómo ajustar un modelo.
La mayor parte de JMP tiene licencia CC-BY-NC, como se encuentra en el archivo LICENSE . Sin embargo, partes del proyecto están disponibles bajo términos de licencia separados:
Si utiliza este código en su investigación, cite el siguiente artículo:
@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}
}
