โปรดดูเว็บไซต์ของเราสำหรับข้อมูลเพิ่มเติมเกี่ยวกับงานนี้
พื้นที่เก็บข้อมูลนี้มีโค้ดสำหรับรายงาน "จากโมเลกุลสู่วัสดุ: การฝึกอบรมแบบจำลองทั่วไปขนาดใหญ่ล่วงหน้าสำหรับการทำนายสมบัติอะตอม"
พื้นที่เก็บข้อมูลนี้เลิกใช้แล้วและไม่ได้รับการดูแลรักษาอีกต่อไป ขณะนี้เรากำลังดำเนินการรวมฟังก์ชันการทำงานของพื้นที่เก็บข้อมูลนี้เข้ากับพื้นที่เก็บข้อมูล Open Catalyst Project อย่างเป็นทางการ หากคุณมีคำถามหรือข้อกังวลใดๆ เกี่ยวกับพื้นที่เก็บข้อมูลนี้ โปรดอย่าลังเลที่จะสร้างปัญหา GitHub หรือติดต่อเราผ่านทางอีเมล กรุณาส่งอีเมลไปที่ Nima Shoghi และ CC Brandon Wood
ในงานนี้ เราขอแนะนำ Joint Multi-domain Pre-training (JMP) ซึ่งเป็นกลยุทธ์ก่อนการฝึกอบรมภายใต้การดูแลที่ฝึกชุดข้อมูลหลายชุดจากโดเมนเคมีที่แตกต่างกันไปพร้อมๆ กัน โดยถือว่าแต่ละชุดข้อมูลเป็นงานก่อนการฝึกอบรมที่ไม่ซ้ำกันภายในกรอบงานแบบหลายงาน . ชุดข้อมูลการฝึกอบรมแบบรวมของเราประกอบด้วยระบบประมาณ 120 ล้านระบบจาก OC20, OC22, ANI-1x และ Transition-1x
ผลงานที่สำคัญของงานนี้คือ:
เราแสดงให้เห็นถึงความสามารถในการวางลักษณะทั่วไปอันทรงพลังของ JMP โดยการประเมินประสิทธิภาพการปรับแต่งอย่างละเอียดในชุดเกณฑ์มาตรฐานที่หลากหลาย ซึ่งครอบคลุมโมเลกุลขนาดเล็ก โมเลกุลขนาดใหญ่ และวัสดุ JMP มีประสิทธิภาพเหนือกว่าการฝึกอบรมตั้งแต่เริ่มต้นอย่างต่อเนื่อง และกำหนดหรือจับคู่เกณฑ์มาตรฐานที่ล้ำสมัยในเกณฑ์มาตรฐานการปรับแต่ง 34 รายการจาก 40 รายการ
เราแสดงให้เห็นว่า JMP ช่วยให้สามารถปรับขนาดโมเดลขนาดใหญ่ได้อย่างมีประสิทธิภาพ ซึ่งปกติแล้วจะพอดีเกินไปหากฝึกตั้งแต่เริ่มต้นบนชุดข้อมูลขนาดเล็ก การฝึกอบรมล่วงหน้าทำหน้าที่เป็นตัวสร้างมาตรฐานที่แข็งแกร่ง ช่วยให้เราสามารถฝึกอบรมโมเดลที่มีพารามิเตอร์ 235M ที่กำหนดประสิทธิภาพที่ล้ำสมัยใหม่ในเกณฑ์มาตรฐานข้อมูลต่ำหลายตัว
เราทำการวิเคราะห์โดยละเอียดเกี่ยวกับข้อกำหนดด้านการคำนวณของ JMP แม้ว่าค่าใช้จ่ายล่วงหน้าจะแพง แต่เราแสดงให้เห็นว่าค่าใช้จ่ายในการฝึกอบรมก่อนการฝึกอบรมของ JMP ได้รับการกู้คืนโดยทำให้การปรับแต่งละเอียดเร็วขึ้นกว่า 12 เท่า เมื่อเทียบกับการฝึกอบรมตั้งแต่เริ่มต้น
ด้วยการฝึกอบรมแบบจำลองขนาดใหญ่ล่วงหน้าเกี่ยวกับข้อมูลทางเคมีที่หลากหลาย เราเชื่อว่า JMP แสดงถึงก้าวสำคัญสู่เป้าหมายของศักยภาพ ML ที่เป็นสากลสำหรับเคมี การเติบโตอย่างต่อเนื่องของข้อมูลที่มีอยู่และพลังการประมวลผลจะช่วยเพิ่มความสามารถของ JMP ในการเรียนรู้การแทนอะตอมที่ถ่ายโอนได้เท่านั้น
JMP แสดงให้เห็นถึงการปรับปรุงโดยเฉลี่ย 59% เมื่อเทียบกับการฝึกอบรมตั้งแต่เริ่มต้น และจับคู่หรือกำหนดงานที่ทันสมัยใน 34 งานจาก 40 งาน งานของเราเน้นย้ำถึงศักยภาพของกลยุทธ์ก่อนการฝึกอบรมที่ใช้ข้อมูลที่หลากหลายเพื่อพัฒนาการคาดการณ์คุณสมบัติในขอบเขตทางเคมี โดยเฉพาะอย่างยิ่งสำหรับงานที่มีข้อมูลน้อย
ขั้นแรก โคลนพื้นที่เก็บข้อมูลและนำทางไปยังไดเร็กทอรีราก:
git clone https://github.com/facebookresearch/JMP.git
cd JMP
จากนั้น ตั้งค่าสภาพแวดล้อม conda ตามที่ระบุในไฟล์ environment.yml
โดยรันคำสั่งต่อไปนี้ (หมายเหตุ: แทนที่ conda
ด้วย 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.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
(ตั้งค่า --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 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
หากต้องการดาวน์โหลดชุดข้อมูล Transition-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 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/
จุดตรวจที่ได้รับการฝึกล่วงหน้าสามารถดาวน์โหลดได้จากลิงค์ต่อไปนี้:
โค้ดเบสของเราได้รับการออกแบบมาให้ใช้ภาษา Python ในการฝึกโมเดล คุณต้องสร้างไฟล์การกำหนดค่าที่ระบุโมเดล ชุดข้อมูล และพารามิเตอร์การฝึก ออบเจ็กต์การกำหนดค่าเหล่านี้ได้รับการตรวจสอบประเภทและตรวจสอบโดยใช้ Pydantic
เมื่อคุณสร้างอ็อบเจ็กต์การกำหนดค่าแล้ว คุณสามารถใช้คลาส jmp.lightning.Runner
เพื่อรวมลูปการฝึกได้ ดู configs/jmp_l_finetune.ipynb
สำหรับตัวอย่างวิธีปรับแต่งโมเดล
JMP ส่วนใหญ่ได้รับใบอนุญาต CC-BY-NC ดังที่พบในไฟล์ LICENSE
อย่างไรก็ตาม บางส่วนของโครงการจะพร้อมใช้งานภายใต้เงื่อนไขสิทธิ์การใช้งานที่แยกต่างหาก:
หากคุณใช้รหัสนี้ในการวิจัยของคุณ โปรดอ้างอิงเอกสารต่อไปนี้:
@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}
}