Super-resolution

In this example, we expand upon the basics of operator learning and demonstrate the use of neural operators for super-resolution.

We will employ the FLAME dataset, a set of flow samples of resolution 32x32 that should be up-sampled to 128x128. You can download the data set from Kaggle and put it into the data/flame directory to reproduce this example.

Setup

import torch
import pathlib
import matplotlib.pyplot as plt
from continuiti.benchmarks.flame import FlameDataset
from continuiti.operators import DeepONet
from continuiti.trainer import Trainer
from continuiti.trainer.callbacks import LearningCurve

Flame Dataset

continuiti provides the Flame class (a special OperatorDataset) that reads and exports samples from the FLAME data. The data set contains train/val splits and has four channels ux, uy, uz, and rho. In this example, we only use channel ux from the first four samples of the val split, and we visualize the provided data using matplotlib.

N = 4
flame_dir = pathlib.Path.cwd().joinpath("..", "data", "flame")
flame = FlameDataset(flame_dir=flame_dir, size=N, split="val", channels=["ux"])

Operator

We define a DeepONet to map the low-resolution data to a continuous function. Note that we increase the expressivity of the trunk network by increasing the width and depth.

operator = DeepONet(shapes=flame.shapes, trunk_width=128, trunk_depth=64)

Training

With an OperatorDataset at hand, training is straightforward using the Trainer.fit method. Here, we add the LearningCurve callback to monitor the training loss.

trainer = Trainer(operator, lr=1e-4)
trainer.fit(flame, epochs=3000, callbacks=[LearningCurve()])

Parameters: 1052272  Device: mps
Epoch 3000/3000  Step 1/1  [====================]  362ms/step  ETA 0:00min - loss/train = 5.1608e-03   

Evaluation

As we can evaluate the trained operator at arbitrary positions, we can plot the mapped function on a fine mesh with 128 positions (or any other resolution), aka super-resolution!

---

Last update: 2024-04-16
Created: 2024-01-15