frameworks:
- pytorch
language:
- en
license: apache-2.0
tags:
- OneScience
- FactFormer
- neural-operator
- factorized-attention
- computational-fluid-dynamics
- Kolmogorov-flow
tasks:
- time-series-prediction
FactFormer
Model Overview
FactFormer (Factorized Transformer) is a Transformer architecture for PDE surrogate modeling. This model predicts the temporal evolution of the vorticity field in two-dimensional Kolmogorov flow, using the first 10 time steps to predict the subsequent 16 by default.
Paper: Scalable Transformer for PDE Surrogate Modeling
Model Description
FactFormer factorizes global attention on a two-dimensional regular grid into attention operations along the two spatial axes. This preserves long-range spatial interactions while reducing computational and memory costs. The model predicts future states in blocks through latent propagation and rolls out the complete forecast window autoregressively.
Use Cases
| Use Case | Description |
|---|---|
| Vorticity time-series prediction | Predict future Kolmogorov-flow states from historical two-dimensional vorticity fields |
| PDE surrogate modeling | Learn mappings from historical to future physical fields on regular grids |
| Factorized-attention research | Evaluate axial attention for high-resolution physical-field modeling |
| Pipeline validation | Validate training and inference using the bundled weights or a small-scale configuration |
Usage
1. OneCode
Use the online OneCode environment for an intelligent, one-click AI for Science (AI4S) programming experience:
Launch OneCode for one-click AI4S programming
2. Manual Setup
Hardware Requirements
- A GPU or DCU is recommended for training and inference.
- A CPU can be used for small-scale pipeline validation, although training with the default configuration will be slow.
- DCU users must install DTK and a PyTorch environment compatible with the target cluster.
Download the Model Package
modelscope download --model OneScience/FactFormer --local_dir ./FactFormer
cd FactFormer
Set Up the Runtime Environment
DCU Environment
# Activate DTK and Conda first
conda create -n onescience311 python=3.11 -y
conda activate onescience311
pip install onescience[cfd-dcu] -i http://mirrors.onescience.ai:3141/pypi/simple/ --trusted-host mirrors.onescience.ai
GPU Environment
# Activate Conda first
conda create -n onescience311 python=3.11 -y libstdcxx-ng=12 libgcc-ng=12 gcc_linux-64=12 gxx_linux-64=12
conda activate onescience311
pip install onescience[cfd-gpu] -i http://mirrors.onescience.ai:3141/pypi/simple/ --trusted-host mirrors.onescience.ai
Training Data
The model uses two-dimensional Kolmogorov-flow data with an original array shape of [120, 320, 256, 256], representing trajectories, time steps, and the two spatial dimensions, respectively. Because the data file is large, it is read on demand using memory mapping.
Download the data with the following command and verify that the data path in conf/config.yaml is configured correctly:
modelscope download --dataset OneScience/Kolmogorov_flow_2d --local_dir ./data
Training
python scripts/train.py
The default weights are saved to weight/factformer_kolmogorov.pt.
Model Weights
This repository provides weights trained on the two-dimensional Kolmogorov-flow dataset in the weight/ directory.
Inference, Evaluation, and Visualization
The model package includes weights for pipeline validation. After preparing the data, run:
python scripts/inference.py
The script loads weight/factformer_kolmogorov.pt by default. Predicted tensors and visualizations are saved to the result/ directory. Training and inference parameters can be modified in conf/config.yaml.
Official OneScience Resources
| Platform | OneScience Repository | Skills Repository |
|---|---|---|
| Gitee | https://gitee.com/onescience-ai/onescience | https://gitee.com/onescience-ai/oneskills |
| GitHub | https://github.com/onescience-ai/OneScience | https://github.com/onescience-ai/oneskills |
Citations and License
- Li, Shu, and Barati Farimani. Scalable Transformer for PDE Surrogate Modeling.
- The Kolmogorov-flow data is used for two-dimensional vorticity prediction at Re=1000.
- This model package is released under the Apache-2.0 license and retains attribution to the original paper and data sources.