Add Transolver model (#591)

* Add transolver model. Signed-off-by: luohk19 <[email protected]> * Add transolver to changelog. Signed-off-by: luohk19 <[email protected]> * Update transolver model's license. Signed-off-by: luohk19 <[email protected]> * add reproduction code of transovler * fix license, add transolver's readme detail * black reformatted * add unitest and fix some bug in train * fix headers * formatting, small modifications * fix: transolver doctest * update transolver * fix: timm.models.layers * formatting * formatting, fix conflicts * fix: doc test * fix: doctest --------- Signed-off-by: luohk19 <[email protected]> Co-authored-by: Mohammad Amin Nabian <[email protected]>
NVIDIA · Oct 17, 2024 · 8876b15 · 8876b15
1 parent f66026a
commit 8876b15
Show file tree

Hide file tree

Showing 18 changed files with 2,281 additions and 0 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -17,6 +17,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Added gradient clipping to StaticCapture utilities.
 - Bistride Multiscale MeshGraphNet example.
 - FIGConvUNet model and example.
+- The Transolver model.
 
 ### Changed
 

diff --git a/docs/img/transolver.png b/docs/img/transolver.png
diff --git a/examples/cfd/darcy_transolver/README.md b/examples/cfd/darcy_transolver/README.md
@@ -0,0 +1,48 @@
+<!-- markdownlint-disable MD033 -->
+
+# Transolver for Darcy Flow
+
+This example demonstrates how to set up a data-driven model for a 2D Darcy flow using
+the Transolver inside of Modulus.
+
+<p align="center">
+<img src="../../../docs/img/transolver.png" />
+</p>
+
+Training progress can be tracked through [MLFlow](https://mlflow.org/docs/latest/index.html).
+This example runs on a single GPU.
+
+## Getting Started
+
+To train the model following modulus's settings, simply run
+
+```bash
+python train_transolver_darcy.py
+```
+
+Each batch is a new data generated by equation, which is different from commonly-used settings.
+
+To reproduce the results in the paper, run
+
+```bash
+python train_transolver_darcy_fix.py
+```
+
+In this case, the train set and test set are fixed after the construction of Dataset,
+corresponding to Transolver's setting.
+
+## Additional Information
+
+In the fixed case, extra data is needed for training and the data path should be added when
+Darcy_2D_fix dataset is constructed. You can download the data
+[here](https://drive.google.com/drive/folders/1UnbQh2WWc6knEHbLn-ZaXrKUZhp7pjt-).
+
+More components are added for convenience. `Validators` calculate the loss between
+ground-truth and prediction, and visualize them in `./mlruns`. Below is a simple example
+of visualization.
+
+[![visualization](https://s21.ax1x.com/2024/09/26/pAlis3T.png)](https://imgse.com/i/pAlis3T)
+
+## References
+
+- [Transolver: A Fast Transformer Solver for PDEs on General Geometries](https://arxiv.org/abs/2402.02366)
diff --git a/examples/cfd/darcy_transolver/config.yaml b/examples/cfd/darcy_transolver/config.yaml
@@ -0,0 +1,63 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: All rights reserved.
+# SPDX-License-Identifier: MIT License
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+model:
+  space_dim: 2
+  n_layers: 3
+  n_hidden: 128
+  dropout: 0.0
+  n_head: 8
+  Time_Input: False
+  act: gelu
+  mlp_ratio: 1
+  fun_dim: 0
+  out_dim: 1
+  slice_dim: 32
+  ref: 8
+  unified_pos: 1
+  slice_num: 32
+
+
+
+normaliser:
+  permeability:
+    mean: 1.25
+    std_dev: .75
+  darcy:
+    mean: 4.52E-2
+    std_dev: 2.79E-2
+
+scheduler:
+  initial_lr: 1.E-3
+  decay_rate: .85
+  decay_pseudo_epochs: 8
+
+training:
+  resolution: 256
+  batch_size: 8
+  rec_results_freq : 8
+  max_pseudo_epochs: 256
+  pseudo_epoch_sample_size: 2048
+
+validation:
+  sample_size: 256
+  validation_pseudo_epochs: 4
diff --git a/examples/cfd/darcy_transolver/config_fix.yaml b/examples/cfd/darcy_transolver/config_fix.yaml
@@ -0,0 +1,64 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: All rights reserved.
+# SPDX-License-Identifier: MIT License
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+model:
+  space_dim: 2
+  n_layers: 8
+  n_hidden: 128
+  dropout: 0.0
+  n_head: 8
+  Time_Input: False
+  act: gelu
+  mlp_ratio: 1
+  fun_dim: 1
+  out_dim: 1
+  slice_dim: 32
+  ref: 8
+  unified_pos: 1
+  slice_num: 64
+
+
+
+normaliser:
+  permeability:
+    mean: 1.25
+    std_dev: .75
+  darcy:
+    mean: 4.52E-2
+    std_dev: 2.79E-2
+
+scheduler:
+  initial_lr: 1.E-3
+  decay_rate: 1.E-5
+  weight_decay: 1.E-5
+  decay_pseudo_epochs: 8
+
+training:
+  resolution: 85
+  batch_size: 4
+  rec_results_freq : 100
+  max_pseudo_epochs: 500
+  pseudo_epoch_sample_size: 1000
+
+validation:
+  sample_size: 200
+  validation_pseudo_epochs: 2