Consulte nosso site para obter mais informações sobre este trabalho.
Este repositório contém o código do artigo "Das moléculas aos materiais: pré-treinamento de grandes modelos generalizáveis para previsão de propriedades atômicas".
Este repositório está obsoleto e não é mais mantido ativamente. Atualmente estamos trabalhando na integração da funcionalidade deste repositório no repositório oficial do Open Catalyst Project. Se você tiver alguma dúvida ou preocupação em relação a este repositório, sinta-se à vontade para criar um problema no github ou entrar em contato conosco por e-mail. Por favor, envie um e-mail para Nima Shoghi e CC Brandon Wood.
Neste trabalho, apresentamos o Joint Multi-domain Pre-training (JMP), uma estratégia de pré-treinamento supervisionado que treina simultaneamente em vários conjuntos de dados de diferentes domínios químicos, tratando cada conjunto de dados como uma tarefa única de pré-treinamento dentro de uma estrutura multitarefa. . Nosso conjunto de dados de treinamento combinado consiste em aproximadamente 120 milhões de sistemas de OC20, OC22, ANI-1x e Transition-1x.
As principais contribuições deste trabalho são:
Demonstramos a poderosa capacidade de generalização do JMP avaliando seu desempenho de ajuste fino em um conjunto diversificado de benchmarks que abrange moléculas pequenas, moléculas grandes e materiais. O JMP supera consistentemente o treinamento desde o início e define ou corresponde ao que há de mais moderno em 34 dos 40 benchmarks de ajuste fino.
Mostramos que o JMP permite o escalonamento eficiente para modelos maiores que normalmente se ajustariam demais se treinados do zero em pequenos conjuntos de dados. O pré-treinamento atua como um forte regularizador, permitindo-nos treinar um modelo com parâmetros de 235M que estabelece um novo desempenho de última geração em vários benchmarks com poucos dados.
Conduzimos uma análise detalhada dos requisitos computacionais do JMP. Embora seja caro no início, mostramos que o custo de pré-treinamento do JMP é recuperado ao permitir um ajuste fino 12 vezes mais rápido em comparação com o treinamento do zero.
Ao pré-treinar grandes modelos em diversos dados químicos, acreditamos que o JMP representa um passo importante em direção ao objetivo de um potencial universal de ML para a química. O crescimento contínuo dos dados disponíveis e do poder computacional só melhorará a capacidade do JMP de aprender representações atômicas transferíveis.
O JMP demonstra uma melhoria média de 59% em relação ao treinamento do zero e combina ou define o que há de mais moderno em 34 das 40 tarefas. Nosso trabalho destaca o potencial de estratégias de pré-treinamento que utilizam dados diversos para avançar na previsão de propriedades em domínios químicos, especialmente para tarefas com poucos dados.
Primeiro, clone o repositório e navegue até o diretório raiz:
git clone https://github.com/facebookresearch/JMP.git
cd JMP Em seguida, configure o ambiente conda, conforme fornecido no arquivo environment.yml . Para fazer isso, execute o seguinte comando (NOTA: substitua conda por mamba se estiver instalado):
conda env create -f environment.yml -n jmpSe o comando acima falhar, você mesmo poderá criar o ambiente manualmente:
# 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 multiprocessEm seguida, ative o ambiente:
conda activate jmpFinalmente, instale o pacote atual da seguinte forma:
pip install -e . O código agora está pronto para ser usado. Consulte os arquivos de configuração no diretório configs para obter exemplos de como executar o código.
Os dados usados para pré-treinamento e ajuste fino não estão incluídos neste repositório devido a restrições de tamanho. No entanto, instruções para download e pré-processamento dos conjuntos de dados OC20, OC22, ANI-1x, Transition-1x, QM9, rMD17, MatBench, QMOF, SPICE e MD22 são fornecidas abaixo.
Para baixar o conjunto de dados ANI-1x e convertê-lo para um formato que possa ser usado por nossa base de código, siga estas etapas:
mkdir -p /path/to/datasets/ani1x
cd /path/to/datasets/ani1x.h5 ) da fonte 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 (defina --num_workers para o número de núcleos de CPU disponíveis): 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 em arquivos .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 baixar o conjunto de dados Transition-1x e convertê-lo para um formato que possa ser usado por nossa base de código, siga estas etapas:
mkdir -p /path/to/datasets/transition1x
cd /path/to/datasets/transition1x.h5 ) da fonte oficial: wget https://figshare.com/ndownloader/files/36035789 -O transition1x-release.h5.traj (defina --num_workers para o número de núcleos de CPU disponíveis): 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 em arquivos .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/ Os pontos de verificação pré-treinados estão disponíveis para download nos seguintes links:
Nossa base de código foi projetada para ser baseada em Python. Para treinar um modelo, você cria um arquivo de configuração que especifica o modelo, o conjunto de dados e os parâmetros de treinamento. Esses objetos de configuração são totalmente verificados e validados usando Pydantic.
Depois de criar um objeto de configuração, você pode usar a classe jmp.lightning.Runner para encerrar o loop de treinamento. Consulte configs/jmp_l_finetune.ipynb para obter um exemplo de como ajustar um modelo.
A maior parte do JMP é licenciada CC-BY-NC, conforme encontrado no arquivo LICENSE . No entanto, partes do projeto estão disponíveis sob termos de licença separados:
Se você usar este código em sua pesquisa, cite o seguinte artigo:
@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}
}
