Insights
AI가 자동으로 큐레이션·번역·정리하는 기술 동향 피드입니다.
EquiformerV3는 효율성과 표현력을 높인 SE(3)-불변 그래프 어텐션 트랜스포머 모델입니다. 이 연구는 OC20 S2EF-2M, MPtrj 등 다양한 데이터셋에서 훈련 및 평가 코드를 제공하며, PyTorch 구현체를 통해 재현성을 강조합니다.
AI 자동 생성 콘텐츠
본 콘텐츠는 GitHub AI Research의 원문을 AI가 자동으로 요약·번역·분석한 것입니다. 원 저작권은 원저작자에게 있으며, 정확한 내용은 반드시 원문을 확인해 주세요.
원문 바로가기이 저장소는 "EquiformerV3: 효율적이고 표현력이 뛰어나며 일반적인 SE(3)-불변 그래프 어텐션 트랜스포머의 확장"이라는 연구의 공식 PyTorch 구현체를 담고 있습니다.
저희는 OC20 S2EF-2M, MPtrj, OMat24, sAlex 데이터셋에서의 훈련 코드와 Matbench Discovery에서의 평가 코드를 제공합니다.
추후 더 많은 데이터셋에 대한 추가적인 훈련 설정이 추가될 예정입니다.
본 저장소는 fairchem의 이 버전을 기반으로 합니다.
재현성을 용이하게 하기 위해 원래의 코드베이스를 포함했으며, 저희 연구와 관련된 코드는 experimental 폴더 아래에 배치했습니다.
환경 설정을 위해서는 여기를 참조하십시오.
OC20 S2EF 데이터셋은 해당 GitHub 저장소의 지침을 따라 다운로드할 수 있습니다.
예를 들어, 다음 명령을 실행하여 OC20 S2EF-2M 데이터셋을 다운로드할 수 있습니다:
cd ocp
python scripts/download_data.py --task s2ef --split "2M" --num-workers 8
이제 흡착 에너지 라벨 대신 총 에너지 라벨로 학습하기 때문에 --ref-energy를 제거했습니다.
데이터셋(MPtrj_2022.9_full.json)을 여기에서 다운로드합니다.
이 파일에서 .json 데이터셋의 경로와 처리된 데이터셋을 저장할 경로를 업데이트해야 합니다 (experimental/datasets/mptrj_convert_json_to_aselmdb.py#L17 및 experimental/datasets/mptrj_convert_json_to_aselmdb.py#L18).
다음 명령을 실행하여 .json 데이터셋을 .aselmdb 데이터셋으로 변환합니다:
python experimental/datasets/mptrj_convert_json_to_aselmdb.py
어떤 원자도 6Å 이내에 이웃이 없는 구조는 제거합니다.
더 나은 로드 밸런싱을 위해 각 구조의 에지(edge) 수를 기록하는 metadata.npz를 생성합니다:
# .aselmdb 파일이 포함된 디렉터리 경로
ASELMDB_DATASET=""
...
### OMat24 및 sAlex
1. 데이터셋은 [여기](https://huggingface.co/datasets/facebook/OMAT24)에서 찾을 수 있습니다.
2. MPtrj와 마찬가지로, 더 나은 로드 밸런싱을 위해 각 구조의 에지 수를 기록하는 `metadata.npz`를 생성합니다:
```bash
# .aselmdb 파일이 포함된 디렉터리 경로
ASELMDB_DATASET=""
...
train 및 val 하위 디렉터리 전체에 대해 2번을 반복합니다. 즉, aimd-from-PBE-1000-npt, aimd-from-PBE-1000-nvt, aimd-from-PBE-3000-npt, aimd-from-PBE-3000-nvt, rattled-1000, rattled-1000-subsampled, rattled-300, rattled-300-subsampled, rattled-500, rattled-500-subsampled, rattled-relax에 대해 이 작업을 수행해야 합니다.저희 작업과 관련된 모든 파일은 experimental 아래에 배치했습니다.
configs에는 학습 및 평가용 설정(config) 파일이 포함되어 있습니다.datasets에는 데이터셋 전처리 유틸리티 함수가 포함되어 있습니다.models에는 EquiformerV3 (+ DeNS) 모델이 포함되어 있습니다.scripts에는 학습 및 평가용 스크립트가 포함되어 있습니다.tasks에는 Matbench Discovery에서 시뮬레이션을 실행하고, 등변성(equivariance)을 테스트하며, 체(body)-순서 실험을 수행하는 코드가 포함되어 있습니다.trainers에는 학습 및 평가용 코드가 포함되어 있습니다.b. 결과 저장 경로를 훈련 스크립트에서 업데이트합니다. 결과 저장 경로.
c. 실행:
```bash
sh experimental/scripts/train/oc20/s2ef/equiformer_v3/equiformer_v3_splits@[email protected]
```
여기에서 $L_{max} = 4$로 EquiformerV3를 훈련하기 위한 설정 파일(config)과 스크립트를 제공합니다.
MPtrj 데이터의 전처리 과정은 여기를 참고하세요.
직접 사전 학습 (Direct pre-training)
a. 설정 파일에서 훈련 세트 경로 (전체 MPtrj 데이터셋)와 검증 세트 경로 (최종 평가는 Matbench Discovery에서 수행되므로 sAlex 검증 세트의 일부를 사용함)를 수정합니다.
c. 훈련 스크립트는 각 노드에 8개의 GPU가 있는 두 개의 노드를 필요로 합니다. 저희는 이 훈련 스크립트가 2개 노드에서 분산 훈련을 시작하는 예시를 제공하며, 훈련은 다양한 방식으로 시작될 수 있음을 알려드립니다.
d. 실행:
```bash
bash experimental/scripts/train/omat24/equiformer_v3/equiformer_v3_mptrj.sh
```
2. 사전 훈련된 체크포인트에서 에너지 헤드 제거
a. 직접적인 사전 훈련 후, 다음 명령을 실행하여 체크포인트에서 에너지 헤드를 제거합니다:
```bash
# 직접 사전 훈련의 체크포인트 경로
CHECKPOINT=""
...
b. 이렇게 하면 새로운 체크포인트(`.../checkpoint_no-energy_block.pt`)가 생성되며, 이는 기울기 미세 조정(gradient fine-tuning) 중 모델 가중치를 초기화하는 데 사용됩니다.
3. 기울기 미세 조정 (Gradient fine-tuning)
a. [훈련 세트 경로](experimental/configs/omat24/mptrj/experiments/gradient/equiformer_v3_grad-finetune_N@7_L@4_attn-hidden@32_rbf@10_max-neighbors@300_attn-grid@14-8_ffn-grid@[email protected]_lr@0-5e-5-epochs@10-bs@64x8-wd@[email protected]#L7) (전체 MPtrj 데이터셋)과 [검증 세트 경로](experimental/configs/omat24/mptrj/experiments/gradient/equiformer_v3_grad-finetune_N@7_L@4_attn-hidden@32_rbf@10_max-neighbors@300_attn-grid@14-8_ffn-grid@[email protected]_lr@0-5e-5-epochs@10-bs@64x8-wd@[email protected]#L75) (최종 평가는 Matbench Discovery에서 수행되므로 sAlex 검증 세트의 일부를 사용했습니다)를 [설정 파일](experimental/configs/omat24/mptrj/experiments/gradient/equiformer_v3_grad-finetune_N@7_L@4_attn-hidden@32_rbf@10_max-neighbors@300_attn-grid@14-8_ffn-grid@[email protected]_lr@0-5e-5-epochs@10-bs@64x8-wd@[email protected])에서 수정합니다.
b. 사전 학습 체크포인트 경로([path to pre-trained checkpoint](experimental/configs/omat24/mptrj/experiments/gradient/equiformer_v3_grad-finetune_N@7_L@4_attn-hidden@32_rbf@10_max-neighbors@300_attn-grid@14-8_ffn-grid@[email protected]_lr@0-5e-5-epochs@10-bs@64x8-wd@[email protected]#L233))를 수정합니다. 이 경로는 2번을 실행하여 얻은 `.../checkpoint_no-energy_block.pt`와 같은 형태여야 합니다.
c. [훈련 스크립트](experimental/scripts/train/omat24/equiformer_v3/equiformer_v3_grad_mptrj.sh)에서 결과 저장 경로([path to save results](experimental/scripts/train/omat24/equiformer_v3/equiformer_v3_grad_mptrj.sh#L16))를 수정합니다.
d. [훈련 스크립트](experimental/scripts/train/omat24/equiformer_v3/equiformer_v3_grad_mptrj.sh)는 각 노드에 8개의 GPU가 있는 두 개의 노드를 요구합니다. 저희는 이 훈련 스크립트가 2개 노드에서 분산 훈련을 시작하는 예시를 제공하며, 훈련은 다양한 방식으로 시작될 수 있음을 알려드립니다.
e. 실행:
```bash
bash experimental/scripts/train/omat24/equiformer_v3/equiformer_v3_grad_mptrj.sh
```
### OMat24 → MPtrj 및 sAlex
여기에서 $L_{max} = 4$와 $L_{max} = 6$을 사용하여 EquiformerV3를 훈련하기 위한 설정과 스크립트를 제공합니다.
MPtrj 데이터의 전처리 과정은 [여기를 참조](#mptrj)하세요.
1. OMat24에 대한 직접적인 사전 훈련
a. [훈련 세트 경로](experimental/configs/omat24/omat24/experiments/direct/equiformer_v3_N@7_L@4_attn-hidden@32_rbf@64_max-neighbors@300_attn-grid@14-8_ffn-grid@14_use-gate-force-head_merge-layer-norm_epochs@4-bs@512-wd@[email protected]@[email protected]@[email protected]@[email protected]_no-amp.yml#L5-L15)와 [검증 세트 경로](experimental/configs/omat24/omat24/experiments/direct/equiformer_v3_N@7_L@4_attn-hidden@32_rbf@64_max-neighbors@300_attn-grid@14-8_ffn-grid@14_use-gate-force-head_merge-layer-norm_epochs@4-bs@512-wd@[email protected]@[email protected]@[email protected]@[email protected]_no-amp.yml#L81)를 [설정 파일](experimental/configs/omat24/omat24/experiments/direct/equiformer_v3_N@7_L@4_attn-hidden@32_rbf@64_max-neighbors@300_attn-grid@14-8_ffn-grid@14_use-gate-force-head_merge-layer-norm_epochs@4-bs@512-wd@[email protected]@[email protected]@[email protected]@[email protected]_no-amp.yml)에서 수정합니다.
b. [훈련 스크립트](experimental/scripts/train/omat24/equiformer_v3/equiformer_v3_omat24.sh)에서 [결과 저장 경로](experimental/scripts/train/omat24/equiformer_v3/equiformer_v3_omat24.sh#L16)를 수정합니다.
c. [훈련 스크립트](experimental/scripts/train/omat24/equiformer_v3/equiformer_v3_omat24.sh)는 노드당 8개의 GPU가 장착된 네 개의 노드가 필요합니다. 훈련 스크립트는 네 개의 노드에서 분산 훈련을 실행하는 예시를 제공하며, 훈련은 다양한 방식으로 시작될 수 있습니다.
d. 다음 명령어를 실행합니다:
```bash
bash experimental/scripts/train/omat24/equiformer_v3/equiformer_v3_omat24.sh
```
e. $L_{max} = 6$인 경우 위 단계를 반복합니다.
2. OMat24에서의 그래디언트 파인튜닝 (Gradient fine-tuning)