thuerey-group/Tadpole

Autoencoders as Foundation Models for 3D PDEs with Online Learning

About

This repository contains pre-trained weights of Tadpole, a foundation model for three-dimensional partial differential equations (PDEs).

Tadpole distinguishes itself from existing PDE foundation models in three key aspects:

(1) Tadpole is pre-trained as an autoencoder to learn the inherent representation of PDE solutions, which is more generalizable than the traditional paradigm of training PDE foundation models directly on the dynamics evolution of PDE solutions.

(2) Tadpole is pretrained with online learning. which utilize a GPU-based solver to generate diverse data distribution without IO or storage bottlenecks induced by 3D PDE data.

(3) Tadpole is able to process multiple downstream tasks, including autoencoding, dynamics prediction, and generative modeling.

For more details, please refer to our research paper:

Installation and Loading Pre-trained Weights

Tadpole can be installed via pip:

pip install git+https://github.com/tum-pbs/Tadpole

The pre-trained weights are named as tadpole_{SIZE}_{TYPE}.safetensors, where {SIZE} can be S, B, or L indicating the model size, and {TYPE} can be encoder or decoder indicating the model type. Weights can be loaded through weight_{TYPE} arguments in Tadpole model classes:

from huggingface_hub import hf_hub_download
ae=TadpoleAutoencoder(
    size="B",
    weight_encoder=hf_hub_download(repo_id="thuerey-group/Tadpole",filename="tadpole_b_encoder.safetensors"),
    # or you can also download the weights from Hugging Face and load it locally
    weight_decoder=hf_hub_download(repo_id="thuerey-group/Tadpole",filename="tadpole_b_decoder.safetensors"),
    )

Please refer to our github repository for more details and tutorials.

Note: Currently, we only provide pre-trained weights for the B-size model; S- and L-size models will be released in the near future. Please stay tuned😊!