docs/mindflow/docs/source_zh_cn/data_driven/burgers_FNO1D.ipynb · MindSpore/docs

Create your Gitee Account

Explore and code with more than 12 million developers，Free private repositories ！：）

Clone or Download

burgers_FNO1D.ipynb 17.00 KB

{
 "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)&emsp;[![下载样例代码](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)&emsp;[![查看源文件](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
}

Copy Edit Raw Blame History

authored 2023-11-16 15:22 . update obs links

马建仓 AI 助手

尝试更多

代码解读

代码找茬

代码优化

https://gitee.com/mindspore/docs.git

git@gitee.com:mindspore/docs.git

mindspore

docs

master

MindSpore/docs

About

Releases

Contributors

Activities

MindSpore/docs .gitee-modal { width: 500px !important; }

About

Releases

The Open Source Evaluation Index is derived from the OSS Compass evaluation system, which evaluates projects around the following three dimensions

Contributors

Activities

Search

MindSpore/docs