{ "cells": [ { "attachments": {}, "cell_type": "markdown", "metadata": { "pycharm": { "name": "#%% md\n" } }, "source": [ "# 基于FNO求解一维Burgers\n", "\n", "[![下载Notebook](https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/website-images/master/resource/_static/logo_notebook.svg)](https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/notebook/master/mindflow/zh_cn/data_driven/mindspore_burgers_FNO1D.ipynb) [![下载样例代码](https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/website-images/master/resource/_static/logo_download_code.svg)](https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/notebook/master/mindflow/zh_cn/data_driven/mindspore_burgers_FNO1D.py) [![查看源文件](https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/website-images/master/resource/_static/logo_source.svg)](https://gitee.com/mindspore/docs/blob/master/docs/mindflow/docs/source_zh_cn/data_driven/burgers_FNO1D.ipynb)" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": { "pycharm": { "name": "#%% md\n" } }, "source": [ "## 概述\n", "\n", "计算流体力学是21世纪流体力学领域的重要技术之一,其通过使用数值方法在计算机中对流体力学的控制方程进行求解,从而实现流动的分析、预测和控制。传统的有限元法(finite element method,FEM)和有限差分法(finite difference method,FDM)常用于复杂的仿真流程(物理建模、网格划分、数值离散、迭代求解等)和较高的计算成本,往往效率低下。因此,借助AI提升流体仿真效率是十分必要的。\n", "\n", "近年来,随着神经网络的迅猛发展,为科学计算提供了新的范式。经典的神经网络是在有限维度的空间进行映射,只能学习与特定离散化相关的解。与经典神经网络不同,傅里叶神经算子(Fourier Neural Operator,FNO)是一种能够学习无限维函数空间映射的新型深度学习架构。该架构可直接学习从任意函数参数到解的映射,用于解决一类偏微分方程的求解问题,具有更强的泛化能力。更多信息可参考[Fourier Neural Operator for Parametric Partial Differential Equations](https://arxiv.org/abs/2010.08895)。\n", "\n", "本案例教程介绍利用傅里叶神经算子的1-d Burgers方程求解方法。" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": { "pycharm": { "name": "#%% md\n" } }, "source": [ "## 伯格斯方程(Burgers' equation)\n", "\n", "一维伯格斯方程(1-d Burgers' equation)是一个非线性偏微分方程,具有广泛应用,包括一维粘性流体流动建模。它的形式如下:\n", "\n", "$$\n", "\\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, 1]\n", "$$\n", "\n", "$$\n", "u(x, 0)=u_0(x), \\quad x \\in(0,1)\n", "$$\n", "\n", "其中$u$表示速度场,$u_0$表示初始条件,$\\nu$表示粘度系数。\n" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ "## 问题描述\n", "\n", "本案例利用Fourier Neural Operator学习初始状态到下一时刻状态的映射,实现一维Burgers'方程的求解:\n", "\n", "$$\n", "u_0 \\mapsto u(\\cdot, 1)\n", "$$" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ "## 技术路径\n", "\n", "MindSpore Flow求解该问题的具体流程如下:\n", "\n", "1. 创建数据集。\n", "2. 构建模型。\n", "3. 优化器与损失函数。\n", "4. 模型训练。" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ "## Fourier Neural Operator\n", "\n", "Fourier Neural Operator模型构架如下图所示。图中$w_0(x)$表示初始涡度,通过Lifting Layer实现输入向量的高维映射,然后将映射结果作为Fourier Layer的输入,进行频域信息的非线性变换,最后由Decoding Layer将变换结果映射至最终的预测结果$w_1(x)$。\n", "\n", "Lifting Layer、Fourier Layer以及Decoding Layer共同组成了Fourier Neural Operator。\n", "\n", "![Fourier Neural Operator模型构架](https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/website-images/master/docs/mindflow/docs/source_zh_cn/data_driven/images/FNO.png)\n", "\n", "Fourier Layer网络结构如下图所示。图中V表示输入向量,上框表示向量经过傅里叶变换后,经过线性变换R,过滤高频信息,然后进行傅里叶逆变换;另一分支经过线性变换W,最后通过激活函数,得到Fourier Layer输出向量。\n", "\n", "![Fourier Layer网络结构](https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/website-images/master/docs/mindflow/docs/source_zh_cn/data_driven/images/FNO-2.png)" ] }, { "cell_type": "code", "execution_count": 1, "metadata": { "pycharm": { "name": "#%%\n" } }, "outputs": [], "source": [ "import os\n", "import time\n", "import numpy as np\n", "import mindspore as ms\n", "\n", "from mindspore.amp import DynamicLossScaler, auto_mixed_precision, all_finite\n", "from mindspore import nn, Tensor, set_seed, ops, data_sink, jit, save_checkpoint\n", "from mindspore import dtype as mstype\n", "from mindflow import FNO1D, RelativeRMSELoss, load_yaml_config, get_warmup_cosine_annealing_lr\n", "from mindflow.pde import UnsteadyFlowWithLoss\n" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ "下述`src`包可以在[applications/data_driven/burgers/fno1d/src](https://gitee.com/mindspore/mindscience/tree/master/MindFlow/applications/data_driven/burgers/fno1d/src)下载。" ] }, { "cell_type": "code", "execution_count": 2, "metadata": { "pycharm": { "name": "#%%\n" } }, "outputs": [], "source": [ "from src.dataset import create_training_dataset\n", "\n", "set_seed(0)\n", "np.random.seed(0)\n", "\n", "ms.set_context(mode=ms.GRAPH_MODE, device_target=\"GPU\", device_id=5)\n", "use_ascend = ms.get_context(attr_key='device_target') == \"Ascend\"" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ "从[config](https://gitee.com/mindspore/mindscience/blob/master/MindFlow/applications/data_driven/burgers/fno1d/configs/fno1d.yaml)中获得模型、数据、优化器的参数。" ] }, { "cell_type": "code", "execution_count": 3, "metadata": { "pycharm": { "name": "#%%\n" } }, "outputs": [], "source": [ "config = load_yaml_config('fno1d.yaml')\n", "data_params = config[\"data\"]\n", "model_params = config[\"model\"]\n", "optimizer_params = config[\"optimizer\"]" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": { "pycharm": { "name": "#%% md\n" } }, "source": [ "## 创建数据集\n", "\n", "下载训练与测试数据集: [data_driven/burgers/dataset](https://download.mindspore.cn/mindscience/mindflow/dataset/applications/data_driven/burgers/dataset/)。\n", "\n", "本案例根据Zongyi Li在 [Fourier Neural Operator for Parametric Partial Differential Equations](https://arxiv.org/pdf/2010.08895.pdf) 一文中对数据集的设置生成训练数据集与测试数据集。具体设置如下:\n", "基于周期性边界,生成满足如下分布的初始条件$u_0(x)$:\n", "\n", "$$\n", "u_0 \\sim \\mu, \\mu=\\mathcal{N}\\left(0,625(-\\Delta+25 I)^{-2}\\right)\n", "$$\n", "\n", "本案例选取粘度系数$\\nu=0.1$,并使用分步法求解方程,其中热方程部分在傅里叶空间中精确求解,然后使用前向欧拉方法求解非线性部分。训练集样本量为1000个,测试集样本量为200个。\n" ] }, { "cell_type": "code", "execution_count": 4, "metadata": { "pycharm": { "name": "#%%\n" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Data preparation finished\n", "input_path: (1000, 1024, 1)\n", "label_path: (1000, 1024)\n" ] } ], "source": [ "# create training dataset\n", "train_dataset = create_training_dataset(data_params, shuffle=True)\n", "\n", "# create test dataset\n", "test_input, test_label = np.load(os.path.join(data_params[\"path\"], \"test/inputs.npy\")), \\\n", " np.load(os.path.join(data_params[\"path\"], \"test/label.npy\"))\n", "test_input = Tensor(np.expand_dims(test_input, -2), mstype.float32)\n", "test_label = Tensor(np.expand_dims(test_label, -2), mstype.float32)" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": { "pycharm": { "name": "#%% md\n" } }, "source": [ "## 构建模型\n", "\n", "网络由1层Lifting layer、1层Decoding layer以及多层Fourier Layer叠加组成:\n", "\n", "- Lifting layer对应样例代码中`FNO1D.fc0`,将输出数据$x$映射至高维;\n", "\n", "- 多层Fourier Layer的叠加对应样例代码中`FNO1D.fno_seq`,本案例采用离散傅里叶变换实现时域与频域的转换;\n", "\n", "- Decoding layer对应代码中`FNO1D.fc1`与`FNO1D.fc2`,获得最终的预测值。\n", "\n", "基于上述网络结构,进行模型初始化,其中模型参数可在[配置文件](https://gitee.com/mindspore/mindscience/blob/master/MindFlow/applications/data_driven/burgers/fno1d/configs/fno1d.yaml)中修改。" ] }, { "cell_type": "code", "execution_count": 5, "metadata": { "pycharm": { "name": "#%%\n" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "name:FNO1D_in_channels:1_out_channels:1_resolution:1024_modes:16_width:64_depth:4\n" ] } ], "source": [ "model = FNO1D(in_channels=model_params[\"in_channels\"],\n", " out_channels=model_params[\"out_channels\"],\n", " resolution=model_params[\"resolution\"],\n", " modes=model_params[\"modes\"],\n", " channels=model_params[\"width\"],\n", " depths=model_params[\"depth\"])\n", "model_params_list = []\n", "for k, v in model_params.items():\n", " model_params_list.append(f\"{k}:{v}\")\n", "model_name = \"_\".join(model_params_list)\n", "print(model_name)" ] }, { "cell_type": "markdown", "metadata": { "pycharm": { "name": "#%% md\n" } }, "source": [ "## 优化器与损失函数\n", "\n", "使用相对均方根误差作为网络训练损失函数:\n" ] }, { "cell_type": "code", "execution_count": 6, "metadata": { "pycharm": { "name": "#%%\n" } }, "outputs": [], "source": [ "steps_per_epoch = train_dataset.get_dataset_size()\n", "lr = get_warmup_cosine_annealing_lr(lr_init=optimizer_params[\"initial_lr\"],\n", " last_epoch=optimizer_params[\"train_epochs\"],\n", " steps_per_epoch=steps_per_epoch,\n", " warmup_epochs=1)\n", "optimizer = nn.Adam(model.trainable_params(), learning_rate=Tensor(lr))\n", "\n", "if use_ascend:\n", " loss_scaler = DynamicLossScaler(1024, 2, 100)\n", " auto_mixed_precision(model, 'O1')\n", "else:\n", " loss_scaler = None" ] }, { "attachments": {}, "cell_type": "markdown", "metadata": { "pycharm": { "name": "#%% md\n" } }, "source": [ "## 模型训练\n", "\n", "使用**MindSpore版本>=2.0.0**,我们可以使用函数编程来训练神经网络。`MindSpore Flow` 为非稳态问题 `UnsteadyFlowWithLoss` 提供了一个训练接口,用于模型训练和评估。\n" ] }, { "cell_type": "code", "execution_count": 7, "metadata": { "pycharm": { "name": "#%%\n" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "./summary/name:FNO1D_in_channels:1_out_channels:1_resolution:1024_modes:16_width:64_depth:4\n", "epoch: 1, time elapsed: 21747.305870056152ms, loss: 2.167046070098877\n", "epoch: 2, time elapsed: 5525.397539138794ms, loss: 0.5935954451560974\n", "epoch: 3, time elapsed: 5459.984540939331ms, loss: 0.7349425554275513\n", "epoch: 4, time elapsed: 4948.82869720459ms, loss: 0.6338694095611572\n", "epoch: 5, time elapsed: 5571.3865756988525ms, loss: 0.3174982964992523\n", "epoch: 6, time elapsed: 5712.041616439819ms, loss: 0.3099440038204193\n", "epoch: 7, time elapsed: 5218.639135360718ms, loss: 0.3117891848087311\n", "epoch: 8, time elapsed: 4819.460153579712ms, loss: 0.1810857653617859\n", "epoch: 9, time elapsed: 4968.810081481934ms, loss: 0.1386510729789734\n", "epoch: 10, time elapsed: 4849.36785697937ms, loss: 0.2102256715297699\n", "================================Start Evaluation================================\n", "mean rms_error: 0.027940063\n", "=================================End Evaluation=================================\n", "...\n", "epoch: 91, time elapsed: 4398.104429244995ms, loss: 0.019643772393465042\n", "epoch: 92, time elapsed: 5479.56109046936ms, loss: 0.0641067773103714\n", "epoch: 93, time elapsed: 5549.5476722717285ms, loss: 0.02199840545654297\n", "epoch: 94, time elapsed: 6238.730907440186ms, loss: 0.024467874318361282\n", "epoch: 95, time elapsed: 5434.457778930664ms, loss: 0.025712188333272934\n", "epoch: 96, time elapsed: 6481.106281280518ms, loss: 0.02247200347483158\n", "epoch: 97, time elapsed: 6303.435325622559ms, loss: 0.026637140661478043\n", "epoch: 98, time elapsed: 5162.56856918335ms, loss: 0.030040305107831955\n", "epoch: 99, time elapsed: 5364.72225189209ms, loss: 0.02589748054742813\n", "epoch: 100, time elapsed: 5902.378797531128ms, loss: 0.028599221259355545\n", "================================Start Evaluation================================\n", "mean rms_error: 0.0037017763\n", "=================================End Evaluation=================================\n" ] } ], "source": [ "problem = UnsteadyFlowWithLoss(model, loss_fn=RelativeRMSELoss(), data_format=\"NHWTC\")\n", "\n", "summary_dir = os.path.join(config[\"summary_dir\"], model_name)\n", "print(summary_dir)\n", "\n", "def forward_fn(data, label):\n", " loss = problem.get_loss(data, label)\n", " return loss\n", "\n", "grad_fn = ms.value_and_grad(forward_fn, None, optimizer.parameters, has_aux=False)\n", "\n", "@jit\n", "def train_step(data, label):\n", " loss, grads = grad_fn(data, label)\n", " if use_ascend:\n", " loss = loss_scaler.unscale(loss)\n", " if all_finite(grads):\n", " grads = loss_scaler.unscale(grads)\n", " loss = ops.depend(loss, optimizer(grads))\n", " return loss\n", "\n", "sink_process = data_sink(train_step, train_dataset, 1)\n", "summary_dir = os.path.join(config[\"summary_dir\"], model_name)\n", "\n", "for epoch in range(1, config[\"epochs\"] + 1):\n", " model.set_train()\n", " local_time_beg = time.time()\n", " for _ in range(steps_per_epoch):\n", " cur_loss = sink_process()\n", " print(\"epoch: {}, time elapsed: {}ms, loss: {}\".format(epoch, (time.time() - local_time_beg) * 1000, cur_loss.asnumpy()))\n", "\n", " if epoch % config['eval_interval'] == 0:\n", " model.set_train(False)\n", " print(\"================================Start Evaluation================================\")\n", " rms_error = problem.get_loss(test_input, test_label)/test_input.shape[0]\n", " print(\"mean rms_error:\", rms_error)\n", " print(\"=================================End Evaluation=================================\")\n", " ckpt_dir = os.path.join(summary_dir, \"ckpt\")\n", " if not os.path.exists(ckpt_dir):\n", " os.makedirs(ckpt_dir)\n", " save_checkpoint(model, os.path.join(ckpt_dir, model_params[\"name\"] + '_epoch' + str(epoch)))" ] } ], "metadata": { "kernelspec": { "display_name": "Python 3.9.0 ('py39')", "language": "python", "name": "python3" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.9.0" }, "vscode": { "interpreter": { "hash": "57ace93c29d9374277a79956c3f1b916d7d9a05468d906842f9921d0d494a29f" } } }, "nbformat": 4, "nbformat_minor": 1 }