FNO算子求解Burgers方程

概述

背景

随着神经网络的迅猛发展，基于神经网络的方法为科学计算提供了新的范式。MindSpore Science 套件为这种新的计算范式提供了多种工具与接口，其中就包括了针对数值求解偏微分方程的算子神经网络。接下来，我们选取流体力学中最经典与简单的案例——求解一维情形下的 Burgers 方程作为示例，来介绍 MindSpore Science 套件的部分接口与使用方式。

经典的神经网络是在有限维度的空间进行映射，只能学习与特定离散化相关的解。与经典神经网络不同，傅里叶神经算子（Fourier Neural Operator，FNO）是一种能够学习无限维函数空间映射的新型深度学习架构。该架构可直接学习从任意函数参数到解的映射，用于解决一类偏微分方程的求解问题，具有更强的泛化能力。更多信息可参考Fourier Neural Operator for Parametric Partial Differential Equations。

问题描述

一维伯格斯方程（1-d Burgers' equation）是一个非线性偏微分方程，具有广泛应用，包括一维粘性流体流动建模。它的形式如下：

\[\begin{split} \partial_t u(x, t)+\partial_x (u^2(x, t)/2)=\nu \partial_{xx} u(x, t), \quad x \in(0,1), t \in(0, T] \\ u(x, 0)=u_0(x), \quad x \in(0,1) \end{split}\]

其中\(u\)表示速度场，\(u_0\)表示初始条件，\(\nu\)表示粘度系数。

本案例利用 Fourier Neural Operator 学习初始状态到下一时刻状态的映射，实现一维Burgers'方程的求解：

\[ u_0 \mapsto u(\cdot, t) \]

技术路径

Fourier Neural Operator模型构架如下图所示。图中\(w_0(x)\)表示初始涡度，通过Lifting Layer实现输入向量的高维映射，然后将映射结果作为Fourier Layer的输入，进行频域信息的非线性变换，最后由Decoding Layer将变换结果映射至最终的预测结果\(w_1(x)\)。

Lifting Layer、Fourier Layer以及Decoding Layer共同组成了Fourier Neural Operator。

Fourier Neural Operator模型构架

Fourier Layer网络结构如下图所示。图中V表示输入向量，上框表示向量经过傅里叶变换后，经过线性变换R，过滤高频信息，然后进行傅里叶逆变换；另一分支经过线性变换W，最后通过激活函数，得到Fourier Layer输出向量。

Fourier Layer网络结构

快速开始

数据集下载地址：data_driven/burgers/。将数据集保存在./dataset 路径下。

关键接口导入

工作目录选择为 mindscience/MindFlow/applications/data_driven/burgers/fno1d，在工作目录下创建 train.py，导入如下关键接口：

from mindspore.amp import DynamicLossScaler, auto_mixed_precision, all_finite
from mindspore import context, nn, Tensor, set_seed, ops, data_sink, jit, save_checkpoint
from mindspore import dtype as mstype
from mindscience import FNO1D, RelativeRMSELoss, load_yaml_config, get_warmup_cosine_annealing_lr
from mindscience.pde import UnsteadyFlowWithLoss

创建数据集

基于周期边界条件，生成满足如下分布的初始条件 \(u_0\)：

\[ u_0 \sim \mathcal{N}\left(0,625(-\Delta + 25I)^{-2}\right). \]

该分布已在 create_training_dataset 中实现，只需按如下格式调用：

from src import create_training_dataset, visual
# create training dataset
train_dataset = create_training_dataset(data_params, model_params, shuffle=True)

# create test dataset
test_input, test_label = np.load(os.path.join(data_params["root_dir"], "test/inputs.npy")), \
                         np.load(os.path.join(data_params["root_dir"], "test/label.npy"))
test_input = Tensor(np.expand_dims(test_input, -2), mstype.float32)
test_label = Tensor(np.expand_dims(test_label, -2), mstype.float32)

构建模型

网络由 1 层 Lifting layer、1 层 Decoding layer 以及多层 Fourier Layer 叠加组成，定义如下：

model = FNO1D(in_channels=model_params["in_channels"],
              out_channels=model_params["out_channels"],
              n_modes=model_params["modes"],
              resolutions=model_params["resolutions"],
              hidden_channels=model_params["hidden_channels"],
              n_layers=model_params["depths"],
              projection_channels=4*model_params["hidden_channels"],
              )
model_params_list = []
for k, v in model_params.items():
    model_params_list.append(f"{k}:{v}")
model_name = "_".join(model_params_list)

优化器

选择 Adam 算法作为优化器，并考虑 learning rate 随迭代次数的衰减，以及混合精度训练：

steps_per_epoch = train_dataset.get_dataset_size()
lr = get_warmup_cosine_annealing_lr(lr_init=optimizer_params["learning_rate"],
last_epoch=optimizer_params["epochs"],
steps_per_epoch=steps_per_epoch,warmup_epochs=1)
optimizer = nn.Adam(model.trainable_params(), learning_rate=Tensor(lr))

if use_ascend:
    from mindspore.amp import DynamicLossScaler, auto_mixed_precision, all_finite
    loss_scaler = DynamicLossScaler(1024, 2, 100)
    auto_mixed_precision(model, 'O3')
else:
    loss_scaler = None

模型训练

使用 MindSpore version >= 2.0.0, 我们可以使用函数式编程来训练神经网络。使用非稳态问题的训练接口 UnsteadyFlowWithLoss ，该接口用于模型训练和评估，损失函数选择为相对均方根误差下的 Burgers 方程残差。

problem = UnsteadyFlowWithLoss(model, loss_fn=RelativeRMSELoss(), data_format="NHWTC")

summary_dir = os.path.join(config["summary"]["summary_dir"], model_name)
print(summary_dir)

def forward_fn(data, label):
    loss = problem.get_loss(data, label)
    if use_ascend:
        loss = loss_scaler.scale(loss)
    return loss

grad_fn = ops.value_and_grad(forward_fn, None, optimizer.parameters, has_aux=False)

@jit
def train_step(data, label):
    loss, grads = grad_fn(data, label)
    if use_ascend:
        loss = loss_scaler.unscale(loss)
        if all_finite(grads):
            grads = loss_scaler.unscale(grads)
    loss = ops.depend(loss, optimizer(grads))
    return loss

sink_process = data_sink(train_step, train_dataset, 1)
ckpt_dir = os.path.join(summary_dir, "ckpt")
if not os.path.exists(ckpt_dir):
    os.makedirs(ckpt_dir)
for epoch in range(1, optimizer_params["epochs"] + 1):
    model.set_train()
    local_time_beg = time.time()
    for _ in range(steps_per_epoch):
        cur_loss = sink_process()
        save_checkpoint(model, os.path.join(ckpt_dir, model_params["name"] + '_epoch' + str(epoch)))

训练方式一：在命令行中调用 `train.py` 脚本

export PYTHONPATH=$(cd ../../../../../ && pwd):$PYTHONPATH
python train.py --config_file_path ./configs/fno1d.yaml --device_target Ascend --device_id 0 --mode GRAPH

其中，

--config_file_path 表示配置文件的路径，默认值 './configs/fno1d.yaml'；

--device_target 表示使用的计算平台类型，可以选择 'Ascend' 或 'GPU'，默认值 'GPU'；

--device_id 表示使用的计算卡编号，可按照实际情况填写，默认值 0；

--mode 表示运行的模式，'GRAPH' 表示静态图模式, 'PYNATIVE' 表示动态图模式。

训练方式二：运行 Jupyter Notebook

您可以使用中文版和英文版 Jupyter Notebook 逐行运行训练和验证代码。

结果展示

下图分别表示在 \(t \in (0,3]\) 上的方程数值解以及不同时间下的切片：

FNO1D Solves Burgers

使用接口简介

接口名称	简介
`FNO1D`	来自 `mindscience.models`，定义一维 FNO 算子神经网络。
`RelativeRMSELoss`	来自 `mindscience.common`，相对 MSE 误差函数。
`get_warmup_cosine_annealing_lr`	来自 `mindscience.common`，带热重启的余弦退火，一种应用广泛的学习率的调整策略。
`UnsteadyFlowWithLoss`	来自 `mindscience.pde`，定义非稳态流问题的损失函数。

核心贡献者

gitee id：liulei277, yezhenghao2023, huangwangwen2025

email: liulei2770919@163.com, yezhenghao@isrc.iscas.ac.cn, wangwen@isrc.iscas.ac.cn