# EquiformerV2: Improved Equivariant Transformer for Scaling to Higher-Degree Representations

This repository contains the official PyTorch implementation of the work "EquiformerV2: Improved Equivariant Transformer for Scaling to Higher-Degree Representations".
We provide the code for training the base model setting on the OC20 S2EF-2M dataset.


<p align="center">
	<img src="fig/equiformer_v2.png" alt="photo not available" width="80%" height="80%">
</p>



## Content ##
0. [Environment Setup](#environment-setup)
0. [Training](#training)
0. [File Structure](#file-structure)
0. [Acknowledgement](#acknowledgement)



## Environment Setup ##


### Environment 

See [here](docs/env_setup.md) for setting up the environment.


### OC20

The dataset for different tasks can be downloaded by following instructions in their [GitHub repository](https://github.com/Open-Catalyst-Project/ocp/blob/main/DATASET.md).

We can download the OC20 S2EF-2M dataset by running:
```
    cd ocp
    python scripts/download_data.py --task s2ef --split "2M" --num-workers 8 --ref-energy
```
We also need to download the `"val_id"` data split to run training.

After downloading, place the datasets under `datasets/oc20/` by using `ln -s`.
Take `s2ef` as an example:
```
    cd datasets
    mkdir oc20
    cd oc20
    ln -s ~/ocp/data/s2ef s2ef
```



## Training ##


### OC20

We train EquiformerV2 on the OC20 S2EF-2M dataset by running:
    
```bash
    sh scripts/train/oc20/s2ef/equiformer_v2s/equiformer_v2s_N@12_L@6_M@2_splits@2M_g@8.sh
```

The above script uses 8 GPUs by default while in the paper, we use 16 GPUs to train each individual model on the S2EF-2M dataset.



## File Structure ##

1. [`nets`](nets) includes code of different network architectures for OC20.
2. [`scripts`](scripts) includes scripts for training models on OC20.
3. [`main_oc20.py`](main_oc20.py) is the code for training, evaluating and running relaxation.
4. [`oc20/trainer`](oc20/trainer) contains code for energy and force trainers.
5. [`oc20/configs`](oc20/configs) contains the config files for S2EF.



## Acknowledgement ##

Our implementation is based on PyTorch, PyG, e3nn, timm and SEGNN.