---
license: apache-2.0
---
## DPOT: Auto-Regressive Denoising Operator Transformer for Large-Scale PDE Pre-Training (ICML'2024)

Code for [paper](https://arxiv.org/pdf/2403.03542) DPOT: Auto-Regressive Denoising Operator Transformer for Large-Scale PDE Pre-Training (ICML'2024). It  pretrains neural operator transformers (from **7M** to **1B**)  on multiple PDE datasets. We will release the pre-trained weights soon.

![fig1](./resources/dpot.jpg)

Our pre-trained DPOT achieves the state-of-the-art performance on multiple PDE datasets and could be used for finetuning on different types of downstream PDE problems.

![fig2](./resources/dpot_result.jpg)



### Pre-trained Model Configuration

We have five pre-trained checkpoints of different sizes.

| Size   | Attention dim | MLP dim | Layers | Heads | Model size |
|--------|---------------|---------|--------|-------|------------|
| Tiny   | 512           | 512     | 4      | 4     | 7M         |
| Small  | 1024          | 1024    | 6      | 8     | 30M        |
| Medium | 1024          | 4096    | 12     | 8     | 122M       |
| Large  | 1536          | 6144    | 24     | 16    | 509M       |
| Huge   | 2048          | 8092    | 27     | 8     | 1.03B      |


#### Loading pre-trained model

Here is an example code of loading pre-trained model.
```
model = DPOTNet(img_size=128, patch_size=8, mixing_type='afno', in_channels=4, in_timesteps=10, out_timesteps=1, out_channels=4, normalize=False, embed_dim=512, modes=32, depth=4, n_blocks=4, mlp_ratio=1, out_layer_dim=32, n_cls=12)
model.load_state_dict(torch.load('model_Ti.pth')['model'])
```


#### Datasets

All datasets are stored using hdf5 format, containing  `data`  field. Some datasets are stored with individual hdf5 files, others are stored within a single hdf5 file.

Download the original file from these sources and preprocess them to `/data` folder.

| Dataset       | Link                                                         |
| ------------- | ------------------------------------------------------------ |
| FNO data      | [Here](https://drive.google.com/drive/folders/1UnbQh2WWc6knEHbLn-ZaXrKUZhp7pjt-) |
| PDEBench data | [Here](https://darus.uni-stuttgart.de/dataset.xhtml?persistentId=doi:10.18419/darus-2986) |
| PDEArena data | [Here](https://microsoft.github.io/pdearena/datadownload/)   |
| CFDbench data | [Here](https://cloud.tsinghua.edu.cn/d/435413b55dea434297d1/) |




### Citation

If you use DPOT in your research, please use the following BibTeX entry.

```
@article{hao2024dpot,
  title={DPOT: Auto-Regressive Denoising Operator Transformer for Large-Scale PDE Pre-Training},
  author={Hao, Zhongkai and Su, Chang and Liu, Songming and Berner, Julius and Ying, Chengyang and Su, Hang and Anandkumar, Anima and Song, Jian and Zhu, Jun},
  journal={arXiv preprint arXiv:2403.03542},
  year={2024}
}
```