Silakan lihat situs web kami untuk informasi lebih lanjut tentang pekerjaan ini.
Repositori ini berisi kode untuk makalah "Dari Molekul ke Bahan: Pra-pelatihan Model Besar yang Dapat Digeneralisasikan untuk Prediksi Sifat Atom".
Repositori ini sudah tidak digunakan lagi dan tidak lagi dipelihara secara aktif. Kami sedang berupaya mengintegrasikan fungsionalitas repositori ini ke dalam repositori resmi Open Catalyst Project. Jika Anda mempunyai pertanyaan atau kekhawatiran sehubungan dengan repositori ini, silakan membuat masalah github atau menghubungi kami melalui email. Silakan kirim email ke Nima Shoghi dan CC Brandon Wood.
Dalam karya ini, kami memperkenalkan Pra-pelatihan Multi-domain Gabungan (JMP), sebuah strategi pra-pelatihan yang diawasi yang secara bersamaan melatih beberapa kumpulan data dari domain kimia yang berbeda, memperlakukan setiap kumpulan data sebagai tugas pra-pelatihan unik dalam kerangka multi-tugas . Kumpulan data pelatihan gabungan kami terdiri dari ~120 juta sistem dari OC20, OC22, ANI-1x, dan Transition-1x.
Kontribusi utama dari pekerjaan ini adalah:
Kami mendemonstrasikan kemampuan generalisasi JMP yang kuat dengan mengevaluasi kinerja penyesuaiannya di beragam rangkaian benchmark yang mencakup molekul kecil, molekul besar, dan material. JMP secara konsisten mengungguli pelatihan dari awal dan menetapkan atau mencocokkan yang canggih pada 34 dari 40 tolok ukur penyesuaian.
Kami menunjukkan bahwa JMP memungkinkan penskalaan yang efisien pada model yang lebih besar yang biasanya akan overfit jika dilatih dari awal pada kumpulan data kecil. Pra-pelatihan bertindak sebagai pengatur yang kuat, memungkinkan kami melatih model dengan 235 juta parameter yang menetapkan performa canggih baru pada berbagai tolok ukur data rendah.
Kami melakukan analisis rinci terhadap persyaratan komputasi JMP. Meskipun mahal di awal, kami menunjukkan bahwa biaya pra-pelatihan JMP dapat dipulihkan dengan memungkinkan penyesuaian 12x lebih cepat dibandingkan dengan pelatihan dari awal.
Dengan melakukan pra-pelatihan terhadap model besar mengenai beragam data kimia, kami yakin JMP mewakili langkah penting menuju tujuan potensi ML universal untuk bidang kimia. Pertumbuhan berkelanjutan dari data yang tersedia dan daya komputasi hanya akan meningkatkan kemampuan JMP untuk mempelajari representasi atom yang dapat ditransfer.
JMP menunjukkan peningkatan rata-rata sebesar 59% dibandingkan pelatihan dari awal, dan mencocokkan atau menetapkan yang canggih pada 34 dari 40 tugas. Pekerjaan kami menyoroti potensi strategi pra-pelatihan yang memanfaatkan beragam data untuk memajukan prediksi properti di seluruh domain kimia, terutama untuk tugas-tugas dengan data rendah.
Pertama, kloning repositori dan navigasikan ke direktori root:
git clone https://github.com/facebookresearch/JMP.git
cd JMP
Kemudian, atur lingkungan conda, seperti yang disediakan dalam file environment.yml
. Untuk melakukannya, jalankan perintah berikut (CATATAN: ganti conda
dengan mamba
jika Anda sudah menginstalnya):
conda env create -f environment.yml -n jmp
Jika perintah di atas gagal, Anda dapat membuat sendiri lingkungan secara manual:
# 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
Kemudian, aktifkan lingkungan:
conda activate jmp
Terakhir, instal paket saat ini sebagai berikut:
pip install -e .
Kode sekarang siap digunakan. Lihat file konfigurasi di direktori configs
untuk contoh cara menjalankan kode.
Data yang digunakan untuk pra-pelatihan dan penyesuaian tidak disertakan dalam repositori ini karena keterbatasan ukuran. Namun, petunjuk untuk mengunduh dan melakukan pra-pemrosesan kumpulan data OC20, OC22, ANI-1x, Transition-1x, QM9, rMD17, MatBench, QMOF, SPICE, dan MD22 disediakan di bawah.
Untuk mendownload dataset ANI-1x dan mengonversinya ke format yang dapat digunakan oleh basis kode kami, ikuti langkah-langkah berikut:
mkdir -p /path/to/datasets/ani1x
cd /path/to/datasets/ani1x
.h5
) dari sumber resmi: wget https://springernature.figshare.com/ndownloader/files/18112775 -O ani1x-release.h5
python -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
(atur --num_workers
ke jumlah inti CPU yang Anda miliki): 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
ke file .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 32
python -m jmp.datasets.scripts.ani1x_preprocess.ani1x_linear_ref linref --src lmdb/train --out_path linref.npz
python -m jmp.datasets.scripts.ani1x_preprocess.ani1x_linear_ref compute_mean_std --src lmdb/train --linref_path linref.npz --out_path mean_std.pkl
Untuk mengunduh kumpulan data Transition-1x dan mengonversinya ke format yang dapat digunakan oleh basis kode kami, ikuti langkah-langkah berikut:
mkdir -p /path/to/datasets/transition1x
cd /path/to/datasets/transition1x
.h5
) dari sumber resmi: wget https://figshare.com/ndownloader/files/36035789 -O transition1x-release.h5
.traj
(atur --num_workers
ke jumlah inti CPU yang Anda miliki): 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
ke file .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 32
python -m jmp.datasets.scripts.transition1x_preprocess.trans1x_linear_ref linref --src lmdb/train --out_path linref.npz
python -m jmp.datasets.scripts.transition1x_preprocess.trans1x_linear_ref compute_mean_std --src lmdb/train --linref_path linref.npz --out_path mean_std.pkl
python -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/
Pos pemeriksaan terlatih tersedia untuk diunduh dari tautan berikut:
Basis kode kami dirancang berbasis Python. Untuk melatih model, Anda membuat file konfigurasi yang menentukan model, himpunan data, dan parameter pelatihan. Objek konfigurasi ini sepenuhnya diperiksa tipenya dan divalidasi menggunakan Pydantic.
Setelah Anda membuat objek konfigurasi, Anda dapat menggunakan kelas jmp.lightning.Runner
untuk menyelesaikan loop pelatihan. Lihat configs/jmp_l_finetune.ipynb
untuk contoh cara menyempurnakan model.
Mayoritas JMP berlisensi CC-BY-NC, seperti yang ditemukan dalam file LICENSE
. Namun, sebagian dari proyek ini tersedia berdasarkan ketentuan lisensi terpisah:
Jika Anda menggunakan kode ini dalam penelitian Anda, harap kutip makalah berikut:
@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}
}