이 작업에 대한 자세한 내용은 당사 웹사이트를 참조하세요.
이 저장소에는 "분자에서 재료까지: 원자 특성 예측을 위한 대형 일반화 모델 사전 훈련" 논문의 코드가 포함되어 있습니다.
이 저장소는 더 이상 사용되지 않으며 더 이상 적극적으로 유지 관리되지 않습니다. 우리는 현재 이 저장소의 기능을 공식 Open Catalyst 프로젝트 저장소에 통합하는 작업을 진행 중입니다. 이 저장소와 관련하여 질문이나 우려 사항이 있는 경우 언제든지 Github 문제를 생성하거나 이메일을 통해 문의해 주세요. Nima Shoghi와 CC Brandon Wood에게 이메일을 보내주세요.
이 작업에서는 서로 다른 화학 영역의 여러 데이터 세트를 동시에 학습하고 각 데이터 세트를 다중 작업 프레임워크 내에서 고유한 사전 학습 작업으로 처리하는 지도 사전 학습 전략인 JMP(Joint Multi-domain Pre-training)를 소개합니다. . 우리의 결합된 훈련 데이터 세트는 OC20, OC22, ANI-1x 및 Transition-1x의 ~120M 시스템으로 구성됩니다.
이 작업의 주요 기여는 다음과 같습니다.
우리는 작은 분자, 큰 분자, 물질을 포괄하는 다양한 벤치마크 제품군에서 JMP의 미세 조정 성능을 평가하여 JMP의 강력한 일반화 능력을 입증합니다. JMP는 지속적으로 처음부터 훈련보다 뛰어난 성능을 발휘하며 미세 조정 벤치마크 40개 중 34개에서 최첨단 수준을 설정하거나 일치시킵니다.
JMP를 사용하면 작은 데이터 세트를 처음부터 훈련할 경우 일반적으로 과적합되는 대규모 모델로 효율적으로 확장할 수 있음을 보여줍니다. 사전 훈련은 강력한 정규화 도구 역할을 하여 여러 저데이터 벤치마크에서 새로운 최첨단 성능을 설정하는 235M 매개변수로 모델을 훈련할 수 있습니다.
우리는 JMP의 계산 요구 사항을 자세히 분석합니다. 초기 비용이 많이 들지만 JMP의 사전 훈련 비용은 처음부터 훈련하는 것보다 12배 이상 빠른 미세 조정을 통해 회복된다는 것을 보여줍니다.
JMP는 다양한 화학 데이터에 대한 대규모 모델을 사전 훈련함으로써 화학의 보편적인 ML 잠재력이라는 목표를 향한 중요한 단계를 나타낸다고 믿습니다. 사용 가능한 데이터와 컴퓨팅 성능이 지속적으로 증가하면 전송 가능한 원자 표현을 학습하는 JMP의 능력이 향상될 뿐입니다.
JMP는 처음부터 훈련하는 것에 비해 평균 59%의 향상을 보여주며 40개 작업 중 34개 작업에 대해 최신 기술을 일치시키거나 설정합니다. 우리의 작업은 다양한 데이터를 활용하여 특히 낮은 데이터 작업의 경우 화학 영역 전반에 걸쳐 속성 예측을 향상시키는 사전 훈련 전략의 잠재력을 강조합니다.
먼저 저장소를 복제하고 루트 디렉터리로 이동합니다.
git clone https://github.com/facebookresearch/JMP.git
cd JMP 그런 다음, environment.yml 파일에 제공된 대로 conda 환경을 설정합니다. 이렇게 하려면 다음 명령을 실행하십시오(참고: conda mamba로 설치한 경우 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 사용 가능한 CPU 코어 수로 설정). 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.pklTransition-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 사용 가능한 CPU 코어 수로 설정). 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의 대부분은 LICENSE 파일에서 볼 수 있듯이 CC-BY-NC 라이선스를 받았습니다. 그러나 프로젝트의 일부는 별도의 라이선스 조건에 따라 제공됩니다.
연구에 이 코드를 사용하는 경우 다음 논문을 인용해 주세요.
@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}
}
