equiformer-v2 0.0.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
|
@@ -0,0 +1,216 @@
|
|
|
1
|
+
Metadata-Version: 2.3
|
|
2
|
+
Name: equiformer_v2
|
|
3
|
+
Version: 0.0.0
|
|
4
|
+
Summary: Equiformer v2
|
|
5
|
+
Requires-Dist: ase>=3.22,<4.0
|
|
6
|
+
Requires-Dist: numpy>=1.26,<2.0
|
|
7
|
+
Requires-Dist: myplots>=0.1.1
|
|
8
|
+
Requires-Dist: fairchem-core<2.0
|
|
9
|
+
Requires-Dist: torch==2.9.1+cu126
|
|
10
|
+
Requires-Dist: scipy<1.12
|
|
11
|
+
Requires-Dist: torchvision==0.24.1+cu126
|
|
12
|
+
Requires-Dist: torchaudio==2.11.0+cu126
|
|
13
|
+
Requires-Dist: torch-scatter==2.1.2+pt29cu126
|
|
14
|
+
Requires-Dist: torch-sparse==0.6.18+pt29cu126
|
|
15
|
+
Requires-Dist: torch-geometric==2.6.1
|
|
16
|
+
Requires-Dist: torch-spline-conv==1.2.2+pt29cu126
|
|
17
|
+
Requires-Dist: torch-cluster==1.6.3+pt29cu126
|
|
18
|
+
Requires-Python: >=3.11
|
|
19
|
+
Project-URL: Repository, https://github.com/daniisler/equiformer_v2/tree/feature-python-package
|
|
20
|
+
Description-Content-Type: text/markdown
|
|
21
|
+
|
|
22
|
+
# EquiformerV2: Improved Equivariant Transformer for Scaling to Higher-Degree Representations
|
|
23
|
+
|
|
24
|
+
**[Paper](https://arxiv.org/abs/2306.12059)** | **[OpenReview](https://openreview.net/forum?id=mCOBKZmrzD)** | **[Poster](docs/equiformer_v2_poster.pdf)**
|
|
25
|
+
|
|
26
|
+
This repository contains the official PyTorch implementation of the work "EquiformerV2: Improved Equivariant Transformer for Scaling to Higher-Degree Representations" (ICLR 2024).
|
|
27
|
+
We provide the code for training the base model setting on the OC20 S2EF-2M and S2EF-All+MD datasets.
|
|
28
|
+
|
|
29
|
+
Additionally, EquiformerV2 has been incorporated into [OCP repository](https://github.com/FAIR-Chem/fairchem/tree/main/src/fairchem/core/models/equiformer_v2) and used in [Open Catalyst demo](https://open-catalyst.metademolab.com/).
|
|
30
|
+
|
|
31
|
+
In our subsequent [work](https://arxiv.org/abs/2403.09549), we find that we can generalize self-supervised learning similar to BERT, which we call **DeNS** (**De**noising **N**on-Equilibrium **S**tructures), to 3D atomistic systems to improve the performance of EquiformerV2 on energy and force predictions.
|
|
32
|
+
Please refer to the [paper](https://arxiv.org/abs/2403.09549) and the [code](https://github.com/atomicarchitects/DeNS) for further details.
|
|
33
|
+
|
|
34
|
+
The next generation, EquiformerV3, is [here](https://github.com/atomicarchitects/equiformer_v3).
|
|
35
|
+
We further improve efficiency, expressivity, and generality and achieve state-of-the-art results on direct and gradient prediction and on tasks requiring higher-order derivatives.
|
|
36
|
+
|
|
37
|
+
|
|
38
|
+
<p align="center">
|
|
39
|
+
<img src="fig/equiformer_v2.png" alt="photo not available" width="98%" height="98%">
|
|
40
|
+
</p>
|
|
41
|
+
|
|
42
|
+
<p align="center">
|
|
43
|
+
<img src="fig/equiformer_v2_speed_accuracy_tradeoffs.png" alt="photo not available" width="98%" height="98%">
|
|
44
|
+
</p>
|
|
45
|
+
|
|
46
|
+
<p align="center">
|
|
47
|
+
<img src="fig/equiformer_v2_oc20_results.png" alt="photo not available" width="98%" height="98%">
|
|
48
|
+
</p>
|
|
49
|
+
|
|
50
|
+
<p align="center">
|
|
51
|
+
<img src="fig/equiformer_v2_adsorbml_results.png" alt="photo not available" width="98%" height="98%">
|
|
52
|
+
</p>
|
|
53
|
+
|
|
54
|
+
<p align="center">
|
|
55
|
+
<img src="fig/equiformer_v2_oc22_results.png" alt="photo not available" width="98%" height="98%">
|
|
56
|
+
</p>
|
|
57
|
+
|
|
58
|
+
|
|
59
|
+
## Content ##
|
|
60
|
+
0. [Environment Setup](#environment-setup)
|
|
61
|
+
0. [Changelog](#changelog)
|
|
62
|
+
0. [Training](#training)
|
|
63
|
+
0. [File Structure](#file-structure)
|
|
64
|
+
0. [Checkpoints](#checkpoints)
|
|
65
|
+
0. [Citation](#citation)
|
|
66
|
+
0. [Acknowledgement](#acknowledgement)
|
|
67
|
+
|
|
68
|
+
|
|
69
|
+
|
|
70
|
+
## Environment Setup ##
|
|
71
|
+
|
|
72
|
+
|
|
73
|
+
### Environment
|
|
74
|
+
|
|
75
|
+
See [here](docs/env_setup.md) for setting up the environment.
|
|
76
|
+
|
|
77
|
+
|
|
78
|
+
### OC20
|
|
79
|
+
|
|
80
|
+
The OC20 S2EF dataset can be downloaded by following instructions in their [GitHub repository](https://github.com/Open-Catalyst-Project/ocp/blob/5a7738f9aa80b1a9a7e0ca15e33938b4d2557edd/DATASET.md#download-and-preprocess-the-dataset).
|
|
81
|
+
|
|
82
|
+
For example, we can download the OC20 S2EF-2M dataset by running:
|
|
83
|
+
```
|
|
84
|
+
cd ocp
|
|
85
|
+
python scripts/download_data.py --task s2ef --split "2M" --num-workers 8 --ref-energy
|
|
86
|
+
```
|
|
87
|
+
We also need to download the `"val_id"` data split to run training.
|
|
88
|
+
|
|
89
|
+
After downloading, place the datasets under `datasets/oc20/` by using `ln -s`:
|
|
90
|
+
```
|
|
91
|
+
cd datasets
|
|
92
|
+
mkdir oc20
|
|
93
|
+
cd oc20
|
|
94
|
+
ln -s ~/ocp/data/s2ef s2ef
|
|
95
|
+
```
|
|
96
|
+
|
|
97
|
+
To train on different splits like All and All+MD, we can follow the same link above to download the datasets.
|
|
98
|
+
|
|
99
|
+
|
|
100
|
+
## Changelog ##
|
|
101
|
+
|
|
102
|
+
|
|
103
|
+
Please refer to [here](docs/changelog.md).
|
|
104
|
+
|
|
105
|
+
|
|
106
|
+
## Training ##
|
|
107
|
+
|
|
108
|
+
|
|
109
|
+
### OC20
|
|
110
|
+
|
|
111
|
+
1. We train EquiformerV2 on the OC20 **S2EF-2M** dataset by running:
|
|
112
|
+
|
|
113
|
+
```bash
|
|
114
|
+
sh scripts/train/oc20/s2ef/equiformer_v2/equiformer_v2_N@12_L@6_M@2_splits@2M_g@multi-nodes.sh
|
|
115
|
+
```
|
|
116
|
+
The above script uses 2 nodes with 8 GPUs on each node.
|
|
117
|
+
|
|
118
|
+
If there is an import error, it is possible that [`ocp/ocpmodels/common/utils.py`](https://github.com/Open-Catalyst-Project/ocp/blob/5a7738f9aa80b1a9a7e0ca15e33938b4d2557edd/ocpmodels/common/utils.py#L329) is not modified.
|
|
119
|
+
Please follow [here](docs/env_setup.md) for details.
|
|
120
|
+
|
|
121
|
+
We can also run training on 8 GPUs on 1 node:
|
|
122
|
+
```bash
|
|
123
|
+
sh scripts/train/oc20/s2ef/equiformer_v2/equiformer_v2_N@12_L@6_M@2_splits@2M_g@8.sh
|
|
124
|
+
```
|
|
125
|
+
|
|
126
|
+
2. We train **EquiformerV2 (153M)** on OC20 **S2EF-All+MD** by running:
|
|
127
|
+
```bash
|
|
128
|
+
sh scripts/train/oc20/s2ef/equiformer_v2/equiformer_v2_N@20_L@6_M@3_splits@all+md_g@multi-nodes.sh
|
|
129
|
+
```
|
|
130
|
+
The above script uses 16 nodes with 8 GPUs on each node.
|
|
131
|
+
|
|
132
|
+
3. We train **EquiformerV2 (31M)** on OC20 **S2EF-All+MD** by running:
|
|
133
|
+
```bash
|
|
134
|
+
sh scripts/train/oc20/s2ef/equiformer_v2/equiformer_v2_N@8_L@4_M@2_splits@all+md_g@multi-nodes.sh
|
|
135
|
+
```
|
|
136
|
+
The above script uses 8 nodes with 8 GPUs on each node.
|
|
137
|
+
|
|
138
|
+
4. We can train EquiformerV2 with **DeNS** (**De**noising **N**on-Equilibrium **S**tructures) as an auxiliary task to further improve the performance on energy and force predictions. Please refer to the [code](https://github.com/atomicarchitects/DeNS) for details.
|
|
139
|
+
|
|
140
|
+
|
|
141
|
+
## File Structure ##
|
|
142
|
+
|
|
143
|
+
1. [`nets`](nets) includes code of different network architectures for OC20.
|
|
144
|
+
2. [`scripts`](scripts) includes scripts for training models on OC20.
|
|
145
|
+
3. [`main_oc20.py`](main_oc20.py) is the code for training, evaluating and running relaxation.
|
|
146
|
+
4. [`oc20/trainer`](oc20/trainer) contains code for the force trainer as well as some utility functions.
|
|
147
|
+
5. [`oc20/configs`](oc20/configs) contains config files for S2EF.
|
|
148
|
+
|
|
149
|
+
|
|
150
|
+
## Checkpoints ##
|
|
151
|
+
|
|
152
|
+
We provide the checkpoints of EquiformerV2 trained on S2EF-2M dataset for 30 epochs, EquiformerV2 (31M) trained on S2EF-All+MD, and EquiformerV2 (153M) trained on S2EF-All+MD.
|
|
153
|
+
|Model |Split |Download |val force MAE (meV / Å) |val energy MAE (meV) |
|
|
154
|
+
|--- |--- |--- |--- |--- |
|
|
155
|
+
|EquiformerV2 |2M |[checkpoint](https://dl.fbaipublicfiles.com/opencatalystproject/models/2023_06/oc20/s2ef/eq2_83M_2M.pt) \| [config](oc20/configs/s2ef/2M/equiformer_v2/equiformer_v2_N@12_L@6_M@2_epochs@30.yml) |19.4 | 278 |
|
|
156
|
+
|EquiformerV2 (31M)|All+MD |[checkpoint](https://dl.fbaipublicfiles.com/opencatalystproject/models/2023_06/oc20/s2ef/eq2_31M_ec4_allmd.pt) \| [config](oc20/configs/s2ef/all_md/equiformer_v2/equiformer_v2_N@8_L@4_M@2_31M.yml) |16.3 | 232 |
|
|
157
|
+
|EquiformerV2 (153M) |All+MD | [checkpoint](https://dl.fbaipublicfiles.com/opencatalystproject/models/2023_06/oc20/s2ef/eq2_153M_ec4_allmd.pt) \| [config](oc20/configs/s2ef/all_md/equiformer_v2/equiformer_v2_N@20_L@6_M@3_153M.yml) |15.0 | 227 |
|
|
158
|
+
|
|
159
|
+
|
|
160
|
+
## Citation ##
|
|
161
|
+
|
|
162
|
+
Please consider citing the works below if this repository is helpful:
|
|
163
|
+
|
|
164
|
+
- [EquiformerV2](https://arxiv.org/abs/2306.12059):
|
|
165
|
+
```bibtex
|
|
166
|
+
@inproceedings{
|
|
167
|
+
equiformer_v2,
|
|
168
|
+
title={{EquiformerV2: Improved Equivariant Transformer for Scaling to Higher-Degree Representations}},
|
|
169
|
+
author={Yi-Lun Liao and Brandon Wood and Abhishek Das* and Tess Smidt*},
|
|
170
|
+
booktitle={International Conference on Learning Representations (ICLR)},
|
|
171
|
+
year={2024},
|
|
172
|
+
url={https://openreview.net/forum?id=mCOBKZmrzD}
|
|
173
|
+
}
|
|
174
|
+
```
|
|
175
|
+
|
|
176
|
+
- [eSCN](https://arxiv.org/abs/2302.03655):
|
|
177
|
+
```bibtex
|
|
178
|
+
@inproceedings{
|
|
179
|
+
escn,
|
|
180
|
+
title={{Reducing SO(3) Convolutions to SO(2) for Efficient Equivariant GNNs}},
|
|
181
|
+
author={Passaro, Saro and Zitnick, C Lawrence},
|
|
182
|
+
booktitle={International Conference on Machine Learning (ICML)},
|
|
183
|
+
year={2023}
|
|
184
|
+
}
|
|
185
|
+
```
|
|
186
|
+
|
|
187
|
+
- [Equiformer](https://arxiv.org/abs/2206.11990):
|
|
188
|
+
```bibtex
|
|
189
|
+
@inproceedings{
|
|
190
|
+
equiformer,
|
|
191
|
+
title={{Equiformer: Equivariant Graph Attention Transformer for 3D Atomistic Graphs}},
|
|
192
|
+
author={Yi-Lun Liao and Tess Smidt},
|
|
193
|
+
booktitle={International Conference on Learning Representations (ICLR)},
|
|
194
|
+
year={2023},
|
|
195
|
+
url={https://openreview.net/forum?id=KwmPfARgOTD}
|
|
196
|
+
}
|
|
197
|
+
```
|
|
198
|
+
|
|
199
|
+
- [OC20 dataset](https://arxiv.org/abs/2010.09990):
|
|
200
|
+
```bibtex
|
|
201
|
+
@article{
|
|
202
|
+
oc20,
|
|
203
|
+
author = {Chanussot*, Lowik and Das*, Abhishek and Goyal*, Siddharth and Lavril*, Thibaut and Shuaibi*, Muhammed and Riviere, Morgane and Tran, Kevin and Heras-Domingo, Javier and Ho, Caleb and Hu, Weihua and Palizhati, Aini and Sriram, Anuroop and Wood, Brandon and Yoon, Junwoong and Parikh, Devi and Zitnick, C. Lawrence and Ulissi, Zachary},
|
|
204
|
+
title = {{Open Catalyst 2020 (OC20) Dataset and Community Challenges}},
|
|
205
|
+
journal = {ACS Catalysis},
|
|
206
|
+
year = {2021},
|
|
207
|
+
doi = {10.1021/acscatal.0c04525},
|
|
208
|
+
}
|
|
209
|
+
```
|
|
210
|
+
|
|
211
|
+
Please direct questions to Yi-Lun Liao (ylliao@mit.edu).
|
|
212
|
+
|
|
213
|
+
|
|
214
|
+
## Acknowledgement ##
|
|
215
|
+
|
|
216
|
+
Our implementation is based on [PyTorch](https://pytorch.org/), [PyG](https://pytorch-geometric.readthedocs.io/en/latest/index.html), [e3nn](https://github.com/e3nn/e3nn), [timm](https://github.com/huggingface/pytorch-image-models), [ocp](https://github.com/Open-Catalyst-Project/ocp), [Equiformer](https://github.com/atomicarchitects/equiformer).
|
|
@@ -0,0 +1,195 @@
|
|
|
1
|
+
# EquiformerV2: Improved Equivariant Transformer for Scaling to Higher-Degree Representations
|
|
2
|
+
|
|
3
|
+
**[Paper](https://arxiv.org/abs/2306.12059)** | **[OpenReview](https://openreview.net/forum?id=mCOBKZmrzD)** | **[Poster](docs/equiformer_v2_poster.pdf)**
|
|
4
|
+
|
|
5
|
+
This repository contains the official PyTorch implementation of the work "EquiformerV2: Improved Equivariant Transformer for Scaling to Higher-Degree Representations" (ICLR 2024).
|
|
6
|
+
We provide the code for training the base model setting on the OC20 S2EF-2M and S2EF-All+MD datasets.
|
|
7
|
+
|
|
8
|
+
Additionally, EquiformerV2 has been incorporated into [OCP repository](https://github.com/FAIR-Chem/fairchem/tree/main/src/fairchem/core/models/equiformer_v2) and used in [Open Catalyst demo](https://open-catalyst.metademolab.com/).
|
|
9
|
+
|
|
10
|
+
In our subsequent [work](https://arxiv.org/abs/2403.09549), we find that we can generalize self-supervised learning similar to BERT, which we call **DeNS** (**De**noising **N**on-Equilibrium **S**tructures), to 3D atomistic systems to improve the performance of EquiformerV2 on energy and force predictions.
|
|
11
|
+
Please refer to the [paper](https://arxiv.org/abs/2403.09549) and the [code](https://github.com/atomicarchitects/DeNS) for further details.
|
|
12
|
+
|
|
13
|
+
The next generation, EquiformerV3, is [here](https://github.com/atomicarchitects/equiformer_v3).
|
|
14
|
+
We further improve efficiency, expressivity, and generality and achieve state-of-the-art results on direct and gradient prediction and on tasks requiring higher-order derivatives.
|
|
15
|
+
|
|
16
|
+
|
|
17
|
+
<p align="center">
|
|
18
|
+
<img src="fig/equiformer_v2.png" alt="photo not available" width="98%" height="98%">
|
|
19
|
+
</p>
|
|
20
|
+
|
|
21
|
+
<p align="center">
|
|
22
|
+
<img src="fig/equiformer_v2_speed_accuracy_tradeoffs.png" alt="photo not available" width="98%" height="98%">
|
|
23
|
+
</p>
|
|
24
|
+
|
|
25
|
+
<p align="center">
|
|
26
|
+
<img src="fig/equiformer_v2_oc20_results.png" alt="photo not available" width="98%" height="98%">
|
|
27
|
+
</p>
|
|
28
|
+
|
|
29
|
+
<p align="center">
|
|
30
|
+
<img src="fig/equiformer_v2_adsorbml_results.png" alt="photo not available" width="98%" height="98%">
|
|
31
|
+
</p>
|
|
32
|
+
|
|
33
|
+
<p align="center">
|
|
34
|
+
<img src="fig/equiformer_v2_oc22_results.png" alt="photo not available" width="98%" height="98%">
|
|
35
|
+
</p>
|
|
36
|
+
|
|
37
|
+
|
|
38
|
+
## Content ##
|
|
39
|
+
0. [Environment Setup](#environment-setup)
|
|
40
|
+
0. [Changelog](#changelog)
|
|
41
|
+
0. [Training](#training)
|
|
42
|
+
0. [File Structure](#file-structure)
|
|
43
|
+
0. [Checkpoints](#checkpoints)
|
|
44
|
+
0. [Citation](#citation)
|
|
45
|
+
0. [Acknowledgement](#acknowledgement)
|
|
46
|
+
|
|
47
|
+
|
|
48
|
+
|
|
49
|
+
## Environment Setup ##
|
|
50
|
+
|
|
51
|
+
|
|
52
|
+
### Environment
|
|
53
|
+
|
|
54
|
+
See [here](docs/env_setup.md) for setting up the environment.
|
|
55
|
+
|
|
56
|
+
|
|
57
|
+
### OC20
|
|
58
|
+
|
|
59
|
+
The OC20 S2EF dataset can be downloaded by following instructions in their [GitHub repository](https://github.com/Open-Catalyst-Project/ocp/blob/5a7738f9aa80b1a9a7e0ca15e33938b4d2557edd/DATASET.md#download-and-preprocess-the-dataset).
|
|
60
|
+
|
|
61
|
+
For example, we can download the OC20 S2EF-2M dataset by running:
|
|
62
|
+
```
|
|
63
|
+
cd ocp
|
|
64
|
+
python scripts/download_data.py --task s2ef --split "2M" --num-workers 8 --ref-energy
|
|
65
|
+
```
|
|
66
|
+
We also need to download the `"val_id"` data split to run training.
|
|
67
|
+
|
|
68
|
+
After downloading, place the datasets under `datasets/oc20/` by using `ln -s`:
|
|
69
|
+
```
|
|
70
|
+
cd datasets
|
|
71
|
+
mkdir oc20
|
|
72
|
+
cd oc20
|
|
73
|
+
ln -s ~/ocp/data/s2ef s2ef
|
|
74
|
+
```
|
|
75
|
+
|
|
76
|
+
To train on different splits like All and All+MD, we can follow the same link above to download the datasets.
|
|
77
|
+
|
|
78
|
+
|
|
79
|
+
## Changelog ##
|
|
80
|
+
|
|
81
|
+
|
|
82
|
+
Please refer to [here](docs/changelog.md).
|
|
83
|
+
|
|
84
|
+
|
|
85
|
+
## Training ##
|
|
86
|
+
|
|
87
|
+
|
|
88
|
+
### OC20
|
|
89
|
+
|
|
90
|
+
1. We train EquiformerV2 on the OC20 **S2EF-2M** dataset by running:
|
|
91
|
+
|
|
92
|
+
```bash
|
|
93
|
+
sh scripts/train/oc20/s2ef/equiformer_v2/equiformer_v2_N@12_L@6_M@2_splits@2M_g@multi-nodes.sh
|
|
94
|
+
```
|
|
95
|
+
The above script uses 2 nodes with 8 GPUs on each node.
|
|
96
|
+
|
|
97
|
+
If there is an import error, it is possible that [`ocp/ocpmodels/common/utils.py`](https://github.com/Open-Catalyst-Project/ocp/blob/5a7738f9aa80b1a9a7e0ca15e33938b4d2557edd/ocpmodels/common/utils.py#L329) is not modified.
|
|
98
|
+
Please follow [here](docs/env_setup.md) for details.
|
|
99
|
+
|
|
100
|
+
We can also run training on 8 GPUs on 1 node:
|
|
101
|
+
```bash
|
|
102
|
+
sh scripts/train/oc20/s2ef/equiformer_v2/equiformer_v2_N@12_L@6_M@2_splits@2M_g@8.sh
|
|
103
|
+
```
|
|
104
|
+
|
|
105
|
+
2. We train **EquiformerV2 (153M)** on OC20 **S2EF-All+MD** by running:
|
|
106
|
+
```bash
|
|
107
|
+
sh scripts/train/oc20/s2ef/equiformer_v2/equiformer_v2_N@20_L@6_M@3_splits@all+md_g@multi-nodes.sh
|
|
108
|
+
```
|
|
109
|
+
The above script uses 16 nodes with 8 GPUs on each node.
|
|
110
|
+
|
|
111
|
+
3. We train **EquiformerV2 (31M)** on OC20 **S2EF-All+MD** by running:
|
|
112
|
+
```bash
|
|
113
|
+
sh scripts/train/oc20/s2ef/equiformer_v2/equiformer_v2_N@8_L@4_M@2_splits@all+md_g@multi-nodes.sh
|
|
114
|
+
```
|
|
115
|
+
The above script uses 8 nodes with 8 GPUs on each node.
|
|
116
|
+
|
|
117
|
+
4. We can train EquiformerV2 with **DeNS** (**De**noising **N**on-Equilibrium **S**tructures) as an auxiliary task to further improve the performance on energy and force predictions. Please refer to the [code](https://github.com/atomicarchitects/DeNS) for details.
|
|
118
|
+
|
|
119
|
+
|
|
120
|
+
## File Structure ##
|
|
121
|
+
|
|
122
|
+
1. [`nets`](nets) includes code of different network architectures for OC20.
|
|
123
|
+
2. [`scripts`](scripts) includes scripts for training models on OC20.
|
|
124
|
+
3. [`main_oc20.py`](main_oc20.py) is the code for training, evaluating and running relaxation.
|
|
125
|
+
4. [`oc20/trainer`](oc20/trainer) contains code for the force trainer as well as some utility functions.
|
|
126
|
+
5. [`oc20/configs`](oc20/configs) contains config files for S2EF.
|
|
127
|
+
|
|
128
|
+
|
|
129
|
+
## Checkpoints ##
|
|
130
|
+
|
|
131
|
+
We provide the checkpoints of EquiformerV2 trained on S2EF-2M dataset for 30 epochs, EquiformerV2 (31M) trained on S2EF-All+MD, and EquiformerV2 (153M) trained on S2EF-All+MD.
|
|
132
|
+
|Model |Split |Download |val force MAE (meV / Å) |val energy MAE (meV) |
|
|
133
|
+
|--- |--- |--- |--- |--- |
|
|
134
|
+
|EquiformerV2 |2M |[checkpoint](https://dl.fbaipublicfiles.com/opencatalystproject/models/2023_06/oc20/s2ef/eq2_83M_2M.pt) \| [config](oc20/configs/s2ef/2M/equiformer_v2/equiformer_v2_N@12_L@6_M@2_epochs@30.yml) |19.4 | 278 |
|
|
135
|
+
|EquiformerV2 (31M)|All+MD |[checkpoint](https://dl.fbaipublicfiles.com/opencatalystproject/models/2023_06/oc20/s2ef/eq2_31M_ec4_allmd.pt) \| [config](oc20/configs/s2ef/all_md/equiformer_v2/equiformer_v2_N@8_L@4_M@2_31M.yml) |16.3 | 232 |
|
|
136
|
+
|EquiformerV2 (153M) |All+MD | [checkpoint](https://dl.fbaipublicfiles.com/opencatalystproject/models/2023_06/oc20/s2ef/eq2_153M_ec4_allmd.pt) \| [config](oc20/configs/s2ef/all_md/equiformer_v2/equiformer_v2_N@20_L@6_M@3_153M.yml) |15.0 | 227 |
|
|
137
|
+
|
|
138
|
+
|
|
139
|
+
## Citation ##
|
|
140
|
+
|
|
141
|
+
Please consider citing the works below if this repository is helpful:
|
|
142
|
+
|
|
143
|
+
- [EquiformerV2](https://arxiv.org/abs/2306.12059):
|
|
144
|
+
```bibtex
|
|
145
|
+
@inproceedings{
|
|
146
|
+
equiformer_v2,
|
|
147
|
+
title={{EquiformerV2: Improved Equivariant Transformer for Scaling to Higher-Degree Representations}},
|
|
148
|
+
author={Yi-Lun Liao and Brandon Wood and Abhishek Das* and Tess Smidt*},
|
|
149
|
+
booktitle={International Conference on Learning Representations (ICLR)},
|
|
150
|
+
year={2024},
|
|
151
|
+
url={https://openreview.net/forum?id=mCOBKZmrzD}
|
|
152
|
+
}
|
|
153
|
+
```
|
|
154
|
+
|
|
155
|
+
- [eSCN](https://arxiv.org/abs/2302.03655):
|
|
156
|
+
```bibtex
|
|
157
|
+
@inproceedings{
|
|
158
|
+
escn,
|
|
159
|
+
title={{Reducing SO(3) Convolutions to SO(2) for Efficient Equivariant GNNs}},
|
|
160
|
+
author={Passaro, Saro and Zitnick, C Lawrence},
|
|
161
|
+
booktitle={International Conference on Machine Learning (ICML)},
|
|
162
|
+
year={2023}
|
|
163
|
+
}
|
|
164
|
+
```
|
|
165
|
+
|
|
166
|
+
- [Equiformer](https://arxiv.org/abs/2206.11990):
|
|
167
|
+
```bibtex
|
|
168
|
+
@inproceedings{
|
|
169
|
+
equiformer,
|
|
170
|
+
title={{Equiformer: Equivariant Graph Attention Transformer for 3D Atomistic Graphs}},
|
|
171
|
+
author={Yi-Lun Liao and Tess Smidt},
|
|
172
|
+
booktitle={International Conference on Learning Representations (ICLR)},
|
|
173
|
+
year={2023},
|
|
174
|
+
url={https://openreview.net/forum?id=KwmPfARgOTD}
|
|
175
|
+
}
|
|
176
|
+
```
|
|
177
|
+
|
|
178
|
+
- [OC20 dataset](https://arxiv.org/abs/2010.09990):
|
|
179
|
+
```bibtex
|
|
180
|
+
@article{
|
|
181
|
+
oc20,
|
|
182
|
+
author = {Chanussot*, Lowik and Das*, Abhishek and Goyal*, Siddharth and Lavril*, Thibaut and Shuaibi*, Muhammed and Riviere, Morgane and Tran, Kevin and Heras-Domingo, Javier and Ho, Caleb and Hu, Weihua and Palizhati, Aini and Sriram, Anuroop and Wood, Brandon and Yoon, Junwoong and Parikh, Devi and Zitnick, C. Lawrence and Ulissi, Zachary},
|
|
183
|
+
title = {{Open Catalyst 2020 (OC20) Dataset and Community Challenges}},
|
|
184
|
+
journal = {ACS Catalysis},
|
|
185
|
+
year = {2021},
|
|
186
|
+
doi = {10.1021/acscatal.0c04525},
|
|
187
|
+
}
|
|
188
|
+
```
|
|
189
|
+
|
|
190
|
+
Please direct questions to Yi-Lun Liao (ylliao@mit.edu).
|
|
191
|
+
|
|
192
|
+
|
|
193
|
+
## Acknowledgement ##
|
|
194
|
+
|
|
195
|
+
Our implementation is based on [PyTorch](https://pytorch.org/), [PyG](https://pytorch-geometric.readthedocs.io/en/latest/index.html), [e3nn](https://github.com/e3nn/e3nn), [timm](https://github.com/huggingface/pytorch-image-models), [ocp](https://github.com/Open-Catalyst-Project/ocp), [Equiformer](https://github.com/atomicarchitects/equiformer).
|
|
@@ -0,0 +1,37 @@
|
|
|
1
|
+
[build-system]
|
|
2
|
+
requires = ["uv_build >= 0.10.10, <0.12.0"]
|
|
3
|
+
build-backend = "uv_build"
|
|
4
|
+
|
|
5
|
+
[project]
|
|
6
|
+
name = "equiformer_v2"
|
|
7
|
+
version = "0.0.0"
|
|
8
|
+
description = "Equiformer v2"
|
|
9
|
+
readme = "README.md"
|
|
10
|
+
requires-python = ">=3.11"
|
|
11
|
+
dynamic = []
|
|
12
|
+
|
|
13
|
+
dependencies = [
|
|
14
|
+
"ase>=3.22,<4.0",
|
|
15
|
+
"numpy>=1.26,<2.0",
|
|
16
|
+
"myplots>=0.1.1",
|
|
17
|
+
"fairchem-core<2.0",
|
|
18
|
+
"torch==2.9.1+cu126",
|
|
19
|
+
"scipy<1.12",
|
|
20
|
+
"torchvision==0.24.1+cu126",
|
|
21
|
+
"torchaudio==2.11.0+cu126",
|
|
22
|
+
"torch_scatter==2.1.2+pt29cu126",
|
|
23
|
+
"torch_sparse==0.6.18+pt29cu126",
|
|
24
|
+
"torch_geometric==2.6.1",
|
|
25
|
+
"torch-spline-conv==1.2.2+pt29cu126",
|
|
26
|
+
"torch-cluster==1.6.3+pt29cu126",
|
|
27
|
+
]
|
|
28
|
+
|
|
29
|
+
[tool.uv]
|
|
30
|
+
find-links = ["https://data.pyg.org/whl/torch-2.9.1+cu126.html"]
|
|
31
|
+
|
|
32
|
+
[[tool.uv.index]]
|
|
33
|
+
name = "pytorch-cu126"
|
|
34
|
+
url = "https://download.pytorch.org/whl/cu126"
|
|
35
|
+
|
|
36
|
+
[project.urls]
|
|
37
|
+
Repository = "https://github.com/daniisler/equiformer_v2/tree/feature-python-package"
|