File size: 2,860 Bytes
81ab294 f8ad77b 81ab294 f8ad77b 81ab294 63c961c f8ad77b 81ab294 f8ad77b 81ab294 63c961c 81ab294 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 |
---
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.
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.
### 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}
}
``` |