diff --git a/MindFlow/applications/research/cascade_net/Cascade-Net.ipynb b/MindFlow/applications/research/cascade_net/Cascade-Net.ipynb new file mode 100644 index 0000000000000000000000000000000000000000..eadb2c40b5b71f13cc6bff8c74752784930b48d9 --- /dev/null +++ b/MindFlow/applications/research/cascade_net/Cascade-Net.ipynb @@ -0,0 +1,645 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": { + "pycharm": { + "name": "#%% md\n" + } + }, + "source": [ + "# Prediction of spatiotemporal field of pulsation velocity in cylindrical wake by Cascade Net\n", + "\n", + "## Environment Setup\n", + "\n", + "This notebook requires **MindSpore version >= 2.0.0** to support new APIs including: *mindspore.jit, mindspore.jit_class, mindspore.data_sink*. Please check [MindSpore Installation](https://www.mindspore.cn/install/en) for details.\n", + "\n", + "In addition, **MindFlow version >=0.1.0** is also required. If it has not been installed in your environment, please select the right version and hardware, then install it as follows." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "mindflow_version = \"0.1.0\" # update if needed\n", + "# GPU Comment out the following code if you are using NPU.\n", + "!pip uninstall -y mindflow-gpu\n", + "!pip install mindflow-gpu==$mindflow_version\n", + "\n", + "# NPU Uncomment if needed.\n", + "# !pip uninstall -y mindflow-ascend\n", + "# !pip install mindflow-ascend==$mindflow_version" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Background\n", + "\n", + "In the process of turbulent spatiotemporal evolution, the pulsating velocity field includes a series of important fluid physical processes, such as separation, transition, and energy transfer. At high Reynolds numbers, the pulsating velocity field exhibits significant nonlinear characteristics. There are vortex structures in turbulent wake that range from maximum to minimum scales, and these fluid motion patterns constitute complex flow field structural characteristics. The process of energy transfer from large-scale structures to small-scale structures in these flow field structures is called the energy cascade physics principle. Inspired by this principle, the small-scale prediction problem can be transformed into a step-by-step prediction problem from large-scale to small-scale.\n", + "\n", + "## Model framework\n", + "\n", + "The model framework is as shown in the following figure:\n", + "\n", + "![Cascade-Net](images/Cascade-Net.png)\n", + "\n", + "Where, the generator is a U-Net structure with spatial and channel attention gates, and its framework is shown in the following figure:\n", + "\n", + "![The U-Net structure of the generator with spatial and channel attention gates](images/The_U-Net_structure_of_the_generator_with_spatial_and_channel_attention_gates.png)\n", + "\n", + "And for the spatial attention gate *S* and channel attention gate *C*, the integration diagrams are shown as follows:\n", + "\n", + "![Spatial attention gate S](images/Spatial_attention_gate_S.png)\n", + "\n", + "![Channel attention gate C](images/Channel_attention_gate_C.png)\n", + "\n", + "## Preparation\n", + "\n", + "Before practice, ensure that MindSpore of suitable version has been correctly installed. If not, you can run the following command:\n", + "\n", + "* [MindSpore installation page](https://www.mindspore.cn/install) Install MindSpore.\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Training\n", + "\n", + "Import code packs." + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "metadata": { + "collapsed": false, + "jupyter": { + "outputs_hidden": false + }, + "pycharm": { + "name": "#%%\n" + }, + "tags": [] + }, + "outputs": [], + "source": [ + "import gc\n", + "import argparse\n", + "\n", + "from mindspore.dataset import GeneratorDataset\n", + "import mindspore as ms\n", + "from mindspore import ops, Tensor\n", + "import mindspore.nn as nn\n", + "import numpy as np\n", + "import matplotlib.pyplot as plt\n", + "import time\n", + "from tqdm import tqdm\n", + "import matplotlib\n", + "from mindflow.utils import load_yaml_config\n", + "from mindspore import context\n", + "from src.Compound_Model import init_sub_model, DefineCompoundCritic, DefineCompoundGan\n", + "from src.Loss import WassersteinLoss, GradLoss\n", + "from src.dataset import AccesstrainDataset, validation_test_dataset\n", + "matplotlib.use('agg') # to prevent \"Fail to create pixmap with Tk_GetPixmap in TkImgPhotoInstanceSetSize\"\n", + "ms.dataset.config.set_prefetch_size(1)\n", + "context.set_context(device_target=\"GPU\", mode=context.PYNATIVE_MODE) \n", + "# context.set_context(device_target=\"Ascend\", mode=context.PYNATIVE_MODE) " + ] + }, + { + "cell_type": "markdown", + "metadata": { + "collapsed": false, + "jupyter": { + "outputs_hidden": false + }, + "pycharm": { + "name": "#%%\n" + } + }, + "source": [ + "## Datasets Preparation\n", + "\n", + "Dataset download link: [Cascade_Net/dataset](https://download-mindspore.osinfra.cn/mindscience/mindflow/dataset/applications/research/Cascade_Net/). Save the dataset under path `./dataset`.\n", + "\n", + "The datasets are mat type file, each containing three dictionaries: train, validation, and test, representing training dataset, validation dataset, and test dataset. The samples number of training dataset is 10792, the samples number of validation datase is 1704, and the samples number of test dataset is 3408. The Reynolds number distribution of training set data, validation set data, and test set data is shown in the figure:\n", + "\n", + "![Sample cases for training validation and testing](images/Sample_cases_for_training_validation_and_testing.png)\n", + "\n", + "- **conditions_Re_pd.mat** including Reynolds number information with a dimensions(*sample*, *Re*) of (, 1), where *Re* is Reynolds number。\n", + "- **Input_flucReal_Cp_pd.mat** including spatiotemporal characteristic matrix of wall pulsating pressure with a dimensions(*sample*, *θ*, *t*, *C*) of (, 128, 128, 3), where *θ* is measurement point phase of wall pressure field, *t* is time where the time length is controled by the characteristic period of pulsating pressure, *C* is the channel number (3 channels are 3 measuring planes of wall pulsating pressure).\n", + "- **Input_scaling_u_pd.mat** including flow field supplementary input featureswith a dimensions (*sample*, *U*) of (, 20), where *U* is supplementary input features of 10 sampling points flow field *u* and *v*.\n", + "- **ur_1.mat**，**ur_3.mat**，**ur_5.mat**，**ur_10.mat** including flow field information at four scales with a dimensions (*sample*, *H*, *W*, *C*) of (, 128, 128, 2) where *H* and *W* are flow field resolution, *C* is the channel number(2 channels are velocity U and V).\n" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "metadata": { + "collapsed": false, + "jupyter": { + "outputs_hidden": false + }, + "pycharm": { + "name": "#%%\n" + } + }, + "outputs": [], + "source": [ + "path = 'dataset/'\n", + "loader = AccesstrainDataset(path)\n", + "train_dataset = GeneratorDataset(source=loader, column_names=[\n", + " 'u_r10_train', 'u_r5_train', 'u_r3_train', 'u_r1_train', 'cp_fluc_train', 'Re_c_train', 'scaling_input_train'])\n", + "train_dataset = train_dataset.shuffle(buffer_size=25)\n", + "\n", + "(u_r10_validation, u_r5_validation, u_r3_validation, u_r1_validation, cp_fluc_validation, Re_c_validation,\n", + " scaling_input_validation, u_r10_test, u_r5_test, u_r3_test, u_r1_test, cp_fluc_test, Re_c_test, scaling_input_test) \\\n", + " = validation_test_dataset(path)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "pycharm": { + "name": "#%% md\n" + } + }, + "source": [ + "## Configure Network and Training Parameters\n", + "\n", + "Load network related parameters from the configuration file." + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "metadata": { + "collapsed": false, + "jupyter": { + "outputs_hidden": false + }, + "pycharm": { + "name": "#%%\n" + } + }, + "outputs": [], + "source": [ + "config_file_path = \"./config/Cascade-Net.yaml\"\n", + "config = load_yaml_config(config_file_path)\n", + "sample_interval = config[\"sample_interval\"]\n", + "batch_size = config[\"batch_size\"]\n", + "train_dataset = train_dataset.batch(batch_size, drop_remainder=True)\n", + "n_channel_p = config[\"n_channel_p\"]\n", + "n_channel_u = config[\"n_channel_u\"]\n", + "lambda_GP = config[\"lambda_GP\"]\n", + "lambda_L2_u = config[\"lambda_L2_u\"]\n", + "lambda_L2_gradu = config[\"lambda_L2_gradu\"]\n", + "critic_model_lr = config[\"critic_model_lr\"]\n", + "gan_model_lr = config[\"gan_model_lr\"]\n", + "n_critic = config[\"n_critic\"]\n", + "sample_num_train = loader.__len__()\n", + "sample_num_val = u_r1_validation.shape[0]\n", + "sample_num_test = u_r1_test.shape[0]\n", + "plot_n = config[\"plot_n\"]\n", + "n_imgs = config[\"n_imgs\"]\n", + "merge_n_imgs = config[\"merge_n_imgs\"]\n", + "latent_z_n_channel = config[\"latent_z_n_channel\"] # z = (N,8,8,50)\n", + "epochs = config[\"epochs\"]" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "pycharm": { + "name": "#%% md\n" + } + }, + "source": [ + "## Model Building\n", + "\n", + "Build sub networks at various scales using the initial_model(), including sub generators and sub discriminators at four scales, and use them in the initialization of generator and discriminator.\n" + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "metadata": {}, + "outputs": [], + "source": [ + "# initialization of generator and discriminator\n", + "(merge_model, g_model_I, d_model_I, g_model_II, d_model_II,\n", + " g_model_III, d_model_III, g_model_IV, d_model_IV) = init_sub_model(n_channel_p, n_channel_u)\n", + "critic_model = DefineCompoundCritic(n_channel_p, n_channel_u, batch_size,\n", + " d_model_I, d_model_II, d_model_III, d_model_IV)\n", + "critic_model.update_parameters_name('critic')\n", + "gan_model = DefineCompoundGan(n_channel_p, n_channel_u, merge_model,\n", + " g_model_I, g_model_II, g_model_III, g_model_IV)\n", + "gan_model.update_parameters_name('generator')" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Train Function\n", + "\n", + "The main part code of model training, including the definition of the loss function, the forward calculation process of the discriminator and generator, and the training process.\n" + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "metadata": {}, + "outputs": [], + "source": [ + "# loss function\n", + "wasserstein_loss = WassersteinLoss()\n", + "mae_loss = nn.MAELoss(reduction='none')\n", + "grad_loss = GradLoss()\n", + "\n", + "# forward calculation of discriminator\n", + "def d_forward_fn(input): \n", + " g_input = [input[4], input[5], input[6], input[7]] \n", + " g_pred = gan_model(g_input)\n", + " d_input = [input[0], input[1], input[2], input[3], g_pred[4], g_pred[0], g_pred[1], g_pred[2], g_pred[3]]\n", + "\n", + " d_pred = critic_model(d_input)\n", + " loss_wass = wasserstein_loss(ops.stack(d_pred[:8]), ops.Concat(axis=0)(\n", + " [-ops.ones_like(ops.stack(d_pred[:4])), ops.ones_like(ops.stack(d_pred[4:8]))]))\n", + " loss_gradp = ops.stack([d_pred[12](), d_pred[13](), d_pred[14](), d_pred[15]()])\n", + " loss = ops.sum(1 * loss_wass) + ops.sum(lambda_GP * loss_gradp)\n", + " loss_list = ms.numpy.concatenate((ms.numpy.expand_dims(loss, 0), loss_wass, loss_gradp), axis=0).asnumpy()\n", + " return loss, loss_list\n", + "\n", + "# forward calculation of generator\n", + "def g_forward_fn(true, input):\n", + " g_input = [input[4], input[5], input[6], input[7]]\n", + " g_pred = gan_model(g_input)\n", + " d_input = [input[0], input[1], input[2], input[3], g_pred[4], g_pred[0], g_pred[1], g_pred[2], g_pred[3]]\n", + " d_pred = critic_model(d_input)\n", + " \n", + " loss_wass = wasserstein_loss(ops.stack(d_pred[4:8]), -ops.ones_like(ops.stack(d_pred[4:8])))\n", + " loss_mae = mae_loss(ops.stack(g_pred[:4]), ops.stack(true[:4]))\n", + " loss_mae = ops.mean(loss_mae, axis=(1, 2, 3, 4))\n", + " loss_grad = grad_loss(ops.stack(g_pred[5:9]), ops.stack(true[:4]))\n", + " loss = (ops.sum(1 * loss_wass) + ops.sum(lambda_L2_u * loss_mae) + ops.sum(lambda_L2_gradu * loss_grad))\n", + " loss_list = ms.numpy.concatenate((ms.numpy.expand_dims(loss, 0), loss_wass, loss_mae, loss_grad), axis=0).asnumpy()\n", + " return loss, loss_list\n", + "\n", + "# training process\n", + "def train_step(g_real_data, input):\n", + " for _ in range(n_critic): \n", + " (_, d_loss_list), d_grads = d_grad_fn(input)\n", + " d_optimizer(d_grads)\n", + " ave_g_loss_train_ncritic.append(d_loss_list)\n", + " ave_d_loss_train.append(np.mean(ave_g_loss_train_ncritic, axis=0))\n", + " (_, g_loss_list), g_grads = g_grad_fn(g_real_data, input)\n", + " g_optimizer(g_grads)\n", + " ave_g_loss_train.append(g_loss_list)\n", + "\n", + "# evaluation process\n", + "def define_evaluation(u_r10, u_r5, u_r3, u_r1, cp, Re, scaling_input, latent_z, idx):\n", + " _, batch_d_loss = d_forward_fn(\n", + " [Tensor(u_r10[idx, :, :, :]), Tensor(u_r5[idx, :, :, :]), Tensor(u_r3[idx, :, :, :]), Tensor(u_r1[idx, :, :, :]),\n", + " Tensor(cp[idx, :, :, :]), Tensor(Re[idx, :]), Tensor(scaling_input[idx, :]), latent_z])\n", + " _, batch_g_loss = g_forward_fn(\n", + " [Tensor(u_r10[idx, :, :, :]), Tensor(u_r5[idx, :, :, :]), Tensor(u_r3[idx, :, :, :]), Tensor(u_r1[idx, :, :, :])],\n", + " [Tensor(u_r10[idx, :, :, :]), Tensor(u_r5[idx, :, :, :]), Tensor(u_r3[idx, :, :, :]), Tensor(u_r1[idx, :, :, :]),\n", + " Tensor(cp[idx, :, :, :]), Tensor(Re[idx, :]), Tensor(scaling_input[idx, :]), latent_z])\n", + " return batch_d_loss, batch_g_loss\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Optimizer\n", + "\n", + "Apply the RMSProp optimizer.\n" + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "metadata": {}, + "outputs": [], + "source": [ + "# optimizer\n", + "d_optimizer = nn.RMSProp(critic_model.trainable_params(), learning_rate=critic_model_lr)\n", + "d_grad_fn = ms.value_and_grad(d_forward_fn, None, d_optimizer.parameters, has_aux=True)\n", + "g_optimizer = nn.RMSProp(gan_model.trainable_params(), learning_rate=gan_model_lr)\n", + "g_grad_fn = ms.value_and_grad(g_forward_fn, None, g_optimizer.parameters, has_aux=True)\n", + "d_optimizer.update_parameters_name('optim_d')\n", + "g_optimizer.update_parameters_name('optim_g')" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Visualization\n", + "\n", + "Visualization of training result and prediction result.\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "y1 = np.linspace(-1, 1, 128)\n", + "x1 = np.linspace(-1, 1, 128)\n", + "xx, yy = np.meshgrid(x1, y1) # xy_data.new\n", + "\n", + "def sample_images(epoch): # plot_n = 5\n", + " idx_train = np.random.randint(0, sample_num_train, plot_n)\n", + " idx_test = np.random.randint(0, sample_num_test, plot_n)\n", + " latent_z_input_train = Tensor(np.random.normal(0, 1, (\n", + " plot_n, latent_z_n_channel, merge_n_imgs, merge_n_imgs)), dtype=ms.float32)\n", + " latent_z_input_test = Tensor(np.random.normal(0, 1, (\n", + " plot_n, latent_z_n_channel, merge_n_imgs, merge_n_imgs)), dtype=ms.float32)\n", + " gen_u_I, gen_u_II, gen_u_III, gen_u_IV, _ = gan_model([\n", + " Tensor(loader[idx_train][4]),\n", + " Tensor(loader[idx_train][5]),\n", + " Tensor(loader[idx_train][6]),\n", + " latent_z_input_train])\n", + " gen_u_I_, gen_u_II_, gen_u_III_, gen_u_IV_, _ = gan_model([\n", + " Tensor(cp_fluc_test[idx_test, :, :, :]),\n", + " Tensor(Re_c_test[idx_test, :]),\n", + " Tensor(scaling_input_test[idx_test, :]),\n", + " latent_z_input_test])\n", + " del _\n", + " plot_images(epoch,\n", + " velocity_part=0,\n", + " gen_u_I=gen_u_I,\n", + " gen_u_II=gen_u_II,\n", + " gen_u_III=gen_u_III,\n", + " gen_u_IV=gen_u_IV,\n", + " sel_index=idx_train,\n", + " is_train=True)\n", + " plot_images(epoch,\n", + " velocity_part=1,\n", + " gen_u_I=gen_u_I,\n", + " gen_u_II=gen_u_II,\n", + " gen_u_III=gen_u_III,\n", + " gen_u_IV=gen_u_IV,\n", + " sel_index=idx_train,\n", + " is_train=True)\n", + " plot_images(epoch,\n", + " velocity_part=0,\n", + " gen_u_I=gen_u_I_,\n", + " gen_u_II=gen_u_II_,\n", + " gen_u_III=gen_u_III_,\n", + " gen_u_IV=gen_u_IV_,\n", + " sel_index=idx_test,\n", + " is_test=True)\n", + " plot_images(epoch,\n", + " velocity_part=1,\n", + " gen_u_I=gen_u_I_,\n", + " gen_u_II=gen_u_II_,\n", + " gen_u_III=gen_u_III_,\n", + " gen_u_IV=gen_u_IV_,\n", + " sel_index=idx_test, is_test=True)\n", + "\n", + "def plot_images(epoch, velocity_part, gen_u_I, gen_u_II, gen_u_III, gen_u_IV, sel_index, is_train=False,\n", + " is_test=False):\n", + " # including ur,vr,gradur, and grad_vr\n", + " # velocity_part = 0 for u and 1 for v\n", + " velocity_name = list(['u', 'v'])\n", + " fig, ax = plt.subplots(5, 8, figsize=(20, 16))\n", + " if is_train:\n", + " index = 0\n", + " for i in range(0, 5, 1):\n", + " ax[i, 0].contourf(xx, yy,\n", + " gen_u_I[index, velocity_part, :, :].reshape(n_imgs, n_imgs).asnumpy(),\n", + " cmap='coolwarm')\n", + " ax[i, 0].axis('off')\n", + " ax[i, 1].contourf(xx, yy, loader[sel_index[index]][0][velocity_part, :, :], cmap='coolwarm')\n", + " ax[i, 1].axis('off')\n", + " ax[i, 2].contourf(xx, yy,\n", + " gen_u_II[index, velocity_part, :, :].reshape(n_imgs, n_imgs).asnumpy(),\n", + " cmap='coolwarm')\n", + " ax[i, 2].axis('off')\n", + " ax[i, 3].contourf(xx, yy, loader[sel_index[index]][1][velocity_part, :, :], cmap='coolwarm')\n", + " ax[i, 3].axis('off')\n", + " ax[i, 4].contourf(xx, yy,\n", + " gen_u_III[index, velocity_part, :, :].reshape(n_imgs, n_imgs).asnumpy(),\n", + " cmap='coolwarm')\n", + " ax[i, 4].axis('off')\n", + " ax[i, 5].contourf(xx, yy, loader[sel_index[index]][2][velocity_part, :, :], cmap='coolwarm')\n", + " ax[i, 5].axis('off')\n", + " ax[i, 6].contourf(xx, yy,\n", + " gen_u_IV[index, velocity_part, :, :].reshape(n_imgs, n_imgs).asnumpy(),\n", + " cmap='coolwarm')\n", + " ax[i, 6].axis('off')\n", + " ax[i, 7].contourf(xx, yy, loader[sel_index[index]][3][velocity_part, :, :], cmap='coolwarm')\n", + " ax[i, 7].axis('off')\n", + "\n", + " index = index + 1\n", + " plt.savefig('training_results/images/imgs_train_' + velocity_name[velocity_part] + '%d.png' % (\n", + " epoch))\n", + " plt.close()\n", + " if is_test:\n", + " index = 0\n", + " for i in range(0, 5, 1):\n", + " ax[i, 0].contourf(xx, yy,\n", + " gen_u_I[index, velocity_part, :, :].reshape(n_imgs, n_imgs).asnumpy(),\n", + " cmap='coolwarm')\n", + " ax[i, 0].axis('off')\n", + " ax[i, 1].contourf(xx, yy,\n", + " u_r10_test[sel_index[index], velocity_part, :, :].reshape(n_imgs, n_imgs),\n", + " cmap='coolwarm')\n", + " ax[i, 1].axis('off')\n", + " ax[i, 2].contourf(xx, yy,\n", + " gen_u_II[index, velocity_part, :, :].reshape(n_imgs, n_imgs).asnumpy(),\n", + " cmap='coolwarm')\n", + " ax[i, 2].axis('off')\n", + " ax[i, 3].contourf(xx, yy,\n", + " u_r5_test[sel_index[index], velocity_part, :, :].reshape(n_imgs, n_imgs),\n", + " cmap='coolwarm')\n", + " ax[i, 3].axis('off')\n", + " ax[i, 4].contourf(xx, yy,\n", + " gen_u_III[index, velocity_part, :, :].reshape(n_imgs, n_imgs).asnumpy(),\n", + " cmap='coolwarm')\n", + " ax[i, 4].axis('off')\n", + " ax[i, 5].contourf(xx, yy,\n", + " u_r3_test[sel_index[index], velocity_part, :, :].reshape(n_imgs, n_imgs),\n", + " cmap='coolwarm')\n", + " ax[i, 5].axis('off')\n", + " ax[i, 6].contourf(xx, yy,\n", + " gen_u_IV[index, velocity_part, :, :].reshape(n_imgs, n_imgs).asnumpy(),\n", + " cmap='coolwarm')\n", + " ax[i, 6].axis('off')\n", + " ax[i, 7].contourf(xx, yy,\n", + " u_r1_test[sel_index[index], velocity_part, :, :].reshape(n_imgs, n_imgs),\n", + " cmap='coolwarm')\n", + " ax[i, 7].axis('off')\n", + " index = index + 1\n", + " plt.savefig('training_results/images/imgs_test_' + velocity_name[velocity_part] + '%d.png' % (\n", + " epoch))\n", + " plt.close()\n", + " plt.cla()\n", + " plt.close('all')" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Model Training\n", + "\n", + "Predict is performed during model training. The accuracy on the training dataset, validation dataset, and test dataset, and the visualization of the predict results are saved every sample_interval epochs. Simultaneously, users could save checkpoint file every sample_interval steps.\n" + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "metadata": { + "collapsed": false, + "jupyter": { + "outputs_hidden": false + }, + "pycharm": { + "name": "#%%\n" + } + }, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + " 0%| | 0/2 [00:04 save generated image samples\n", + " if epoch % sample_interval == 0:\n", + " sample_images(epoch)\n", + " ms.save_checkpoint(gan_model, \"training_results/model/gan_%d\" % epoch + \".ckpt\")\n", + " ms.save_checkpoint(critic_model, \"training_results/model/critic_%d\" % epoch + \".ckpt\")\n", + " np.savetxt('training_results/loss/loss.txt', losslog, fmt='%.4f')\n", + " del ave_d_loss_train, ave_g_loss_train, batch_d_loss_val, batch_g_loss_val, batch_d_loss_test, batch_g_loss_test\n", + " gc.collect()" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "pycharm": { + "name": "#%% md\n" + } + }, + "source": [ + "## **Result visualization**\n", + "\n", + "The predicted and true values of the velocity fields at various scales in the U and V directions generated by the network in the train dataset under 300 epochs of training are shown in the following figure:\n", + "\n", + " \n", + "\n", + "The predicted and true values of the velocity fields at various scales in the U and V directions generated by the network in the test dataset under 300 epochs of training are shown in the following figure:\n", + "\n", + " \n", + "\n", + "Here, each row give one sample. And one sample give eight figures, each group has two figures including generated flow field at left and the true flow field at right, and the scale decreases from left to right among each group.\n" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "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.13" + }, + "vscode": { + "interpreter": { + "hash": "b9063439a3781aed32d6b0dd4804a0c8b51ecec7893a0f31b99846bc91ef39eb" + } + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} diff --git a/MindFlow/applications/research/cascade_net/Cascade-Net_CN.ipynb b/MindFlow/applications/research/cascade_net/Cascade-Net_CN.ipynb new file mode 100644 index 0000000000000000000000000000000000000000..a08c9bd6700e6cdb10cefd22c52e40f8796506a9 --- /dev/null +++ b/MindFlow/applications/research/cascade_net/Cascade-Net_CN.ipynb @@ -0,0 +1,630 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": { + "pycharm": { + "name": "#%% md\n" + } + }, + "source": [ + "# Cascade-Net预测圆柱尾迹脉动速度时空场\n", + "\n", + "\n", + "## 环境安装\n", + "\n", + "本案例要求 **MindSpore >= 2.0.0** 版本以调用如下接口: *mindspore.jit, mindspore.jit_class, mindspore.data_sink*。具体请查看[MindSpore安装](https://www.mindspore.cn/install)。\n", + "\n", + "此外，你需要安装 **MindFlow >=0.1.0** 版本。如果当前环境还没有安装，请按照下列方式选择后端和版本进行安装。" + ] + }, + { + "cell_type": "raw", + "metadata": {}, + "source": [ + "mindflow_version = \"0.1.0\" # update if needed\n", + "# GPU Comment out the following code if you are using NPU.\n", + "!pip uninstall -y mindflow-gpu\n", + "!pip install mindflow-gpu==$mindflow_version\n", + "\n", + "# NPU Uncomment if needed.\n", + "# !pip uninstall -y mindflow-ascend\n", + "# !pip install mindflow-ascend==$mindflow_version" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 背景介绍\n", + "\n", + "在湍流时空演化过程中，脉动速度场包含了一系列重要的流体物理过程，如分离、转捩和能量传递。在高雷诺数下，脉动速度场表现出明显的非线性特征。湍流尾迹中存在着从最大尺度到最小尺度的涡结构，这些流体运动模式构成了复杂的流场结构特征。在这些流场结构中，能量从大尺度结构向小尺度结构转移的过程被称为能量级串物理原理。受到这一原理的启发，可以将小尺度预测问题转化为由大尺度向小尺度逐级预测问题。\n", + "\n", + "## 模型框架\n", + "\n", + "模型框架如下图所示：\n", + "\n", + "![Cascade-Net](images/Cascade-Net.png)\n", + "\n", + "图中，Generater为具有空间和通道注意力门的U-Net结构，其框架如下图所示：\n", + "\n", + "![The U-Net structure of the generator with spatial and channel attention gates](images/The_U-Net_structure_of_the_generator_with_spatial_and_channel_attention_gates.png)\n", + "\n", + "其中空间注意力门*S*和通道注意力门*C*示意图如下：\n", + "\n", + "![Spatial attention gate S](images/Spatial_attention_gate_S.png)\n", + "\n", + "![Channel attention gate C](images/Channel_attention_gate_C.png)\n", + "\n", + "## 准备环节\n", + "\n", + "实践前，确保已经正确安装合适版本的MindSpore。如果没有，可以通过：\n", + "\n", + "* [MindSpore安装页面](https://www.mindspore.cn/install) 安装MindSpore。\n", + "\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 训练环节\n", + "\n", + "引入代码包。" + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "metadata": { + "collapsed": false, + "jupyter": { + "outputs_hidden": false + }, + "pycharm": { + "name": "#%%\n" + }, + "tags": [] + }, + "outputs": [], + "source": [ + "import gc\n", + "import argparse\n", + "\n", + "from mindspore.dataset import GeneratorDataset\n", + "import mindspore as ms\n", + "from mindspore import ops, Tensor\n", + "import mindspore.nn as nn\n", + "import numpy as np\n", + "import matplotlib.pyplot as plt\n", + "import time\n", + "from tqdm import tqdm\n", + "import matplotlib\n", + "from mindflow.utils import load_yaml_config\n", + "from mindspore import context\n", + "from src.Compound_Model import init_sub_model, DefineCompoundCritic, DefineCompoundGan\n", + "from src.Loss import WassersteinLoss, GradLoss\n", + "from src.dataset import AccesstrainDataset, validation_test_dataset\n", + "matplotlib.use('agg') # to prevent \"Fail to create pixmap with Tk_GetPixmap in TkImgPhotoInstanceSetSize\"\n", + "ms.dataset.config.set_prefetch_size(1)\n", + "context.set_context(device_target=\"GPU\", mode=context.PYNATIVE_MODE) \n", + "# context.set_context(device_target=\"Ascend\", mode=context.PYNATIVE_MODE) " + ] + }, + { + "cell_type": "markdown", + "metadata": { + "pycharm": { + "name": "#%% md\n" + } + }, + "source": [ + "## 数据集的准备\n", + "\n", + "数据集下载地址：[Cascade_Net/dataset](https://download-mindspore.osinfra.cn/mindscience/mindflow/dataset/applications/research/Cascade_Net/)。将数据集保存在`./dataset`路径下。\n", + "\n", + "数据为mat类型文件，每个文件中包括train、validation、test三个词典，分别表示训练集数据、验证集数据、测试集数据，其中训练集样本10792，验证集样本1704，测试集3408。训练集数据、验证集数据、测试集数据数据雷诺数分布如图所示：\n", + "\n", + "![Sample cases for training validation and testing](images/Sample_cases_for_training_validation_and_testing.png)\n", + "\n", + "- conditions_Re_pd.mat文件包括样本雷诺数信息，其维度(*sample*, *Re*)为(, 1)。其中，*sample*为样本，*Re*为样本雷诺数。\n", + "- Input_flucReal_Cp_pd.mat文件包括样本壁面脉动压力的时空特征矩阵，其维度(*sample*, *θ*, *t*, *C*)为(, 128, 128, 3)。其中，*sample*为样本，*θ*为壁面压力场测点相位，*t*为时间，由脉动压力特征周期确定时间长度，*C*为通道数，3个通道分别为三个测面的压力数据。\n", + "- Input_scaling_u_pd.mat文件包括样本流场补充输入特征，其维度(*sample*, *U*)为(, 20)。其中，*sample*为样本，*U*为10个采样点流场补充输入特征*u*和*v*。\n", + "- ur_1.mat，ur_3.mat，ur_5.mat，ur_10.mat文件为各尺度流场信息，其维度(*sample*, *H*, *W*, *C*)为(, 128, 128, 2)。其中，*sample*为样本，*H*和*W*为流场分辨率，*C*为通道数，2个通道分别为速度*u*和*v*。\n" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "metadata": {}, + "outputs": [], + "source": [ + "path = 'dataset/'\n", + "loader = AccesstrainDataset(path)\n", + "train_dataset = GeneratorDataset(source=loader, column_names=[\n", + " 'u_r10_train', 'u_r5_train', 'u_r3_train', 'u_r1_train', 'cp_fluc_train', 'Re_c_train', 'scaling_input_train'])\n", + "train_dataset = train_dataset.shuffle(buffer_size=25)\n", + "\n", + "(u_r10_validation, u_r5_validation, u_r3_validation, u_r1_validation, cp_fluc_validation, Re_c_validation,\n", + " scaling_input_validation, u_r10_test, u_r5_test, u_r3_test, u_r1_test, cp_fluc_test, Re_c_test, scaling_input_test) \\\n", + " = validation_test_dataset(path)" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "pycharm": { + "name": "#%% md\n" + } + }, + "source": [ + "## 配置网络与训练参数\n", + "\n", + "从配置文件中读取网络相关参数。" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "metadata": { + "collapsed": false, + "jupyter": { + "outputs_hidden": false + }, + "pycharm": { + "name": "#%%\n" + } + }, + "outputs": [], + "source": [ + "config_file_path = \"./config/Cascade-Net.yaml\"\n", + "config = load_yaml_config(config_file_path)\n", + "sample_interval = config[\"sample_interval\"]\n", + "batch_size = config[\"batch_size\"]\n", + "train_dataset = train_dataset.batch(batch_size, drop_remainder=True)\n", + "n_channel_p = config[\"n_channel_p\"]\n", + "n_channel_u = config[\"n_channel_u\"]\n", + "lambda_GP = config[\"lambda_GP\"]\n", + "lambda_L2_u = config[\"lambda_L2_u\"]\n", + "lambda_L2_gradu = config[\"lambda_L2_gradu\"]\n", + "critic_model_lr = config[\"critic_model_lr\"]\n", + "gan_model_lr = config[\"gan_model_lr\"]\n", + "n_critic = config[\"n_critic\"]\n", + "sample_num_train = loader.__len__()\n", + "sample_num_val = u_r1_validation.shape[0]\n", + "sample_num_test = u_r1_test.shape[0]\n", + "plot_n = config[\"plot_n\"]\n", + "n_imgs = config[\"n_imgs\"]\n", + "merge_n_imgs = config[\"merge_n_imgs\"]\n", + "latent_z_n_channel = config[\"latent_z_n_channel\"] # z = (N,8,8,50)\n", + "epochs = config[\"epochs\"]" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 模型构建\n", + "\n", + "通过initial_model()函数构建各尺度子网络，包括四个尺度的子生成器和子判别器，并用于生成器和判别器初始化。\n" + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "metadata": {}, + "outputs": [], + "source": [ + "# 生成器及判别器\n", + "(merge_model, g_model_I, d_model_I, g_model_II, d_model_II,\n", + " g_model_III, d_model_III, g_model_IV, d_model_IV) = init_sub_model(n_channel_p, n_channel_u)\n", + "critic_model = DefineCompoundCritic(n_channel_p, n_channel_u, batch_size,\n", + " d_model_I, d_model_II, d_model_III, d_model_IV)\n", + "critic_model.update_parameters_name('critic')\n", + "gan_model = DefineCompoundGan(n_channel_p, n_channel_u, merge_model,\n", + " g_model_I, g_model_II, g_model_III, g_model_IV)\n", + "gan_model.update_parameters_name('generator')" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 训练函数\n", + "\n", + "训练主体代码部分，包括损失函数的定义、判别器和生成器前向计算过程和训练过程。" + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "metadata": {}, + "outputs": [], + "source": [ + "# 损失函数\n", + "wasserstein_loss = WassersteinLoss()\n", + "mae_loss = nn.MAELoss(reduction='none')\n", + "grad_loss = GradLoss()\n", + "\n", + "# 判别器前向计算过程\n", + "def d_forward_fn(input): \n", + " g_input = [input[4], input[5], input[6], input[7]] \n", + " g_pred = gan_model(g_input)\n", + " d_input = [input[0], input[1], input[2], input[3], g_pred[4], g_pred[0], g_pred[1], g_pred[2], g_pred[3]]\n", + "\n", + " d_pred = critic_model(d_input)\n", + " loss_wass = wasserstein_loss(ops.stack(d_pred[:8]), ops.Concat(axis=0)(\n", + " [-ops.ones_like(ops.stack(d_pred[:4])), ops.ones_like(ops.stack(d_pred[4:8]))]))\n", + " loss_gradp = ops.stack([d_pred[12](), d_pred[13](), d_pred[14](), d_pred[15]()])\n", + " loss = ops.sum(1 * loss_wass) + ops.sum(lambda_GP * loss_gradp)\n", + " loss_list = ms.numpy.concatenate((ms.numpy.expand_dims(loss, 0), loss_wass, loss_gradp), axis=0).asnumpy()\n", + " return loss, loss_list\n", + "\n", + "# 生成器前向计算过程\n", + "def g_forward_fn(true, input):\n", + " g_input = [input[4], input[5], input[6], input[7]]\n", + " g_pred = gan_model(g_input)\n", + " d_input = [input[0], input[1], input[2], input[3], g_pred[4], g_pred[0], g_pred[1], g_pred[2], g_pred[3]]\n", + " d_pred = critic_model(d_input)\n", + " \n", + " loss_wass = wasserstein_loss(ops.stack(d_pred[4:8]), -ops.ones_like(ops.stack(d_pred[4:8])))\n", + " loss_mae = mae_loss(ops.stack(g_pred[:4]), ops.stack(true[:4]))\n", + " loss_mae = ops.mean(loss_mae, axis=(1, 2, 3, 4))\n", + " loss_grad = grad_loss(ops.stack(g_pred[5:9]), ops.stack(true[:4]))\n", + " loss = (ops.sum(1 * loss_wass) + ops.sum(lambda_L2_u * loss_mae) + ops.sum(lambda_L2_gradu * loss_grad))\n", + " loss_list = ms.numpy.concatenate((ms.numpy.expand_dims(loss, 0), loss_wass, loss_mae, loss_grad), axis=0).asnumpy()\n", + " return loss, loss_list\n", + "\n", + "# 训练过程\n", + "def train_step(g_real_data, input):\n", + " # 计算判别器损失和梯度 \n", + " for _ in range(n_critic): \n", + " (_, d_loss_list), d_grads = d_grad_fn(input)\n", + " d_optimizer(d_grads)\n", + " ave_g_loss_train_ncritic.append(d_loss_list)\n", + " ave_d_loss_train.append(np.mean(ave_g_loss_train_ncritic, axis=0))\n", + " (_, g_loss_list), g_grads = g_grad_fn(g_real_data, input)\n", + " g_optimizer(g_grads)\n", + " ave_g_loss_train.append(g_loss_list)\n", + "\n", + "# 测试评价\n", + "def define_evaluation(u_r10, u_r5, u_r3, u_r1, cp, Re, scaling_input, latent_z, idx):\n", + " _, batch_d_loss = d_forward_fn(\n", + " [Tensor(u_r10[idx, :, :, :]), Tensor(u_r5[idx, :, :, :]), Tensor(u_r3[idx, :, :, :]), Tensor(u_r1[idx, :, :, :]),\n", + " Tensor(cp[idx, :, :, :]), Tensor(Re[idx, :]), Tensor(scaling_input[idx, :]), latent_z])\n", + " _, batch_g_loss = g_forward_fn(\n", + " [Tensor(u_r10[idx, :, :, :]), Tensor(u_r5[idx, :, :, :]), Tensor(u_r3[idx, :, :, :]), Tensor(u_r1[idx, :, :, :])],\n", + " [Tensor(u_r10[idx, :, :, :]), Tensor(u_r5[idx, :, :, :]), Tensor(u_r3[idx, :, :, :]), Tensor(u_r1[idx, :, :, :]),\n", + " Tensor(cp[idx, :, :, :]), Tensor(Re[idx, :]), Tensor(scaling_input[idx, :]), latent_z])\n", + " return batch_d_loss, batch_g_loss\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 优化器\n", + "\n", + "当前案例中的优化器采用RMSProp优化器。" + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "metadata": {}, + "outputs": [], + "source": [ + "# 优化器\n", + "d_optimizer = nn.RMSProp(critic_model.trainable_params(), learning_rate=critic_model_lr)\n", + "d_grad_fn = ms.value_and_grad(d_forward_fn, None, d_optimizer.parameters, has_aux=True)\n", + "g_optimizer = nn.RMSProp(gan_model.trainable_params(), learning_rate=gan_model_lr)\n", + "g_grad_fn = ms.value_and_grad(g_forward_fn, None, g_optimizer.parameters, has_aux=True)\n", + "d_optimizer.update_parameters_name('optim_d')\n", + "g_optimizer.update_parameters_name('optim_g')" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 结果可视化\n", + "\n", + "绘制网络的预测图，用于结果输出。" + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "metadata": {}, + "outputs": [], + "source": [ + "y1 = np.linspace(-1, 1, 128)\n", + "x1 = np.linspace(-1, 1, 128)\n", + "xx, yy = np.meshgrid(x1, y1) # xy_data.new\n", + "\n", + "def sample_images(epoch): # plot_n = 5\n", + " idx_train = np.random.randint(0, sample_num_train, plot_n)\n", + " idx_test = np.random.randint(0, sample_num_test, plot_n)\n", + " latent_z_input_train = Tensor(np.random.normal(0, 1, (\n", + " plot_n, latent_z_n_channel, merge_n_imgs, merge_n_imgs)), dtype=ms.float32)\n", + " latent_z_input_test = Tensor(np.random.normal(0, 1, (\n", + " plot_n, latent_z_n_channel, merge_n_imgs, merge_n_imgs)), dtype=ms.float32)\n", + " gen_u_I, gen_u_II, gen_u_III, gen_u_IV, _ = gan_model([\n", + " Tensor(loader[idx_train][4]),\n", + " Tensor(loader[idx_train][5]),\n", + " Tensor(loader[idx_train][6]),\n", + " latent_z_input_train])\n", + " gen_u_I_, gen_u_II_, gen_u_III_, gen_u_IV_, _ = gan_model([\n", + " Tensor(cp_fluc_test[idx_test, :, :, :]),\n", + " Tensor(Re_c_test[idx_test, :]),\n", + " Tensor(scaling_input_test[idx_test, :]),\n", + " latent_z_input_test])\n", + " del _\n", + " plot_images(epoch,\n", + " velocity_part=0,\n", + " gen_u_I=gen_u_I,\n", + " gen_u_II=gen_u_II,\n", + " gen_u_III=gen_u_III,\n", + " gen_u_IV=gen_u_IV,\n", + " sel_index=idx_train,\n", + " is_train=True)\n", + " plot_images(epoch,\n", + " velocity_part=1,\n", + " gen_u_I=gen_u_I,\n", + " gen_u_II=gen_u_II,\n", + " gen_u_III=gen_u_III,\n", + " gen_u_IV=gen_u_IV,\n", + " sel_index=idx_train,\n", + " is_train=True)\n", + " plot_images(epoch,\n", + " velocity_part=0,\n", + " gen_u_I=gen_u_I_,\n", + " gen_u_II=gen_u_II_,\n", + " gen_u_III=gen_u_III_,\n", + " gen_u_IV=gen_u_IV_,\n", + " sel_index=idx_test,\n", + " is_test=True)\n", + " plot_images(epoch,\n", + " velocity_part=1,\n", + " gen_u_I=gen_u_I_,\n", + " gen_u_II=gen_u_II_,\n", + " gen_u_III=gen_u_III_,\n", + " gen_u_IV=gen_u_IV_,\n", + " sel_index=idx_test, is_test=True)\n", + "\n", + "def plot_images(epoch, velocity_part, gen_u_I, gen_u_II, gen_u_III, gen_u_IV, sel_index, is_train=False,\n", + " is_test=False):\n", + " # including ur,vr,gradur, and grad_vr\n", + " # velocity_part = 0 for u and 1 for v\n", + " velocity_name = list(['u', 'v'])\n", + " fig, ax = plt.subplots(5, 8, figsize=(20, 16))\n", + " if is_train:\n", + " index = 0\n", + " for i in range(0, 5, 1):\n", + " ax[i, 0].contourf(xx, yy,\n", + " gen_u_I[index, velocity_part, :, :].reshape(n_imgs, n_imgs).asnumpy(),\n", + " cmap='coolwarm')\n", + " ax[i, 0].axis('off')\n", + " ax[i, 1].contourf(xx, yy, loader[sel_index[index]][0][velocity_part, :, :], cmap='coolwarm')\n", + " ax[i, 1].axis('off')\n", + " ax[i, 2].contourf(xx, yy,\n", + " gen_u_II[index, velocity_part, :, :].reshape(n_imgs, n_imgs).asnumpy(),\n", + " cmap='coolwarm')\n", + " ax[i, 2].axis('off')\n", + " ax[i, 3].contourf(xx, yy, loader[sel_index[index]][1][velocity_part, :, :], cmap='coolwarm')\n", + " ax[i, 3].axis('off')\n", + " ax[i, 4].contourf(xx, yy,\n", + " gen_u_III[index, velocity_part, :, :].reshape(n_imgs, n_imgs).asnumpy(),\n", + " cmap='coolwarm')\n", + " ax[i, 4].axis('off')\n", + " ax[i, 5].contourf(xx, yy, loader[sel_index[index]][2][velocity_part, :, :], cmap='coolwarm')\n", + " ax[i, 5].axis('off')\n", + " ax[i, 6].contourf(xx, yy,\n", + " gen_u_IV[index, velocity_part, :, :].reshape(n_imgs, n_imgs).asnumpy(),\n", + " cmap='coolwarm')\n", + " ax[i, 6].axis('off')\n", + " ax[i, 7].contourf(xx, yy, loader[sel_index[index]][3][velocity_part, :, :], cmap='coolwarm')\n", + " ax[i, 7].axis('off')\n", + "\n", + " index = index + 1\n", + " plt.savefig('training_results/images/imgs_train_' + velocity_name[velocity_part] + '%d.png' % (\n", + " epoch))\n", + " plt.close()\n", + " if is_test:\n", + " index = 0\n", + " for i in range(0, 5, 1):\n", + " ax[i, 0].contourf(xx, yy,\n", + " gen_u_I[index, velocity_part, :, :].reshape(n_imgs, n_imgs).asnumpy(),\n", + " cmap='coolwarm')\n", + " ax[i, 0].axis('off')\n", + " ax[i, 1].contourf(xx, yy,\n", + " u_r10_test[sel_index[index], velocity_part, :, :].reshape(n_imgs, n_imgs),\n", + " cmap='coolwarm')\n", + " ax[i, 1].axis('off')\n", + " ax[i, 2].contourf(xx, yy,\n", + " gen_u_II[index, velocity_part, :, :].reshape(n_imgs, n_imgs).asnumpy(),\n", + " cmap='coolwarm')\n", + " ax[i, 2].axis('off')\n", + " ax[i, 3].contourf(xx, yy,\n", + " u_r5_test[sel_index[index], velocity_part, :, :].reshape(n_imgs, n_imgs),\n", + " cmap='coolwarm')\n", + " ax[i, 3].axis('off')\n", + " ax[i, 4].contourf(xx, yy,\n", + " gen_u_III[index, velocity_part, :, :].reshape(n_imgs, n_imgs).asnumpy(),\n", + " cmap='coolwarm')\n", + " ax[i, 4].axis('off')\n", + " ax[i, 5].contourf(xx, yy,\n", + " u_r3_test[sel_index[index], velocity_part, :, :].reshape(n_imgs, n_imgs),\n", + " cmap='coolwarm')\n", + " ax[i, 5].axis('off')\n", + " ax[i, 6].contourf(xx, yy,\n", + " gen_u_IV[index, velocity_part, :, :].reshape(n_imgs, n_imgs).asnumpy(),\n", + " cmap='coolwarm')\n", + " ax[i, 6].axis('off')\n", + " ax[i, 7].contourf(xx, yy,\n", + " u_r1_test[sel_index[index], velocity_part, :, :].reshape(n_imgs, n_imgs),\n", + " cmap='coolwarm')\n", + " ax[i, 7].axis('off')\n", + " index = index + 1\n", + " plt.savefig('training_results/images/imgs_test_' + velocity_name[velocity_part] + '%d.png' % (\n", + " epoch))\n", + " plt.close()\n", + " plt.cla()\n", + " plt.close('all')" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## 模型训练\n", + "\n", + "模型训练过程中边训练边预测，每训练sample_interval个epoch后输出一次训练集、验证集及测试集上的推理精度并保存可视化结果。同时，还可以每隔sample_interval保存一次checkpoint文件并保存loss文件。" + ] + }, + { + "cell_type": "code", + "execution_count": 9, + "metadata": { + "collapsed": false, + "jupyter": { + "outputs_hidden": false + }, + "pycharm": { + "name": "#%%\n" + } + }, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 2/2 [02:05<00:00, 62.74s/it]\n", + "\n", + "KeyboardInterrupt\n", + "\n" + ] + } + ], + "source": [ + "critic_model.set_train()\n", + "gan_model.set_train()\n", + "\n", + "for epoch in range(epochs):\n", + " start = time.perf_counter()\n", + " ave_d_loss_train = list()\n", + " ave_g_loss_train = list()\n", + " for _, (u_r10_train, u_r5_train, u_r3_train, u_r1_train, cp_fluc_train, Re_c_train, scaling_input_train) in (\n", + " enumerate(tqdm(train_dataset))):\n", + " ave_g_loss_train_ncritic = list()\n", + " latent_z_input_train = ops.normal(\n", + " (batch_size, latent_z_n_channel, merge_n_imgs, merge_n_imgs),\n", + " 0, 1, )\n", + " input = [u_r10_train, u_r5_train, u_r3_train, u_r1_train,\n", + " cp_fluc_train, Re_c_train, scaling_input_train, latent_z_input_train]\n", + " g_true = [u_r10_train, u_r5_train, u_r3_train, u_r1_train]\n", + " train_step(g_true, input)\n", + " ave_d_loss_train = np.mean(ave_d_loss_train, axis=0)\n", + " ave_g_loss_train = np.mean(ave_g_loss_train, axis=0)\n", + "\n", + " # predict the d_loss and g_loss in validation\n", + " idx_ = np.random.randint(0, sample_num_val, batch_size)\n", + " latent_z_input_validation = Tensor(np.random.normal(0, 1, (\n", + " batch_size, latent_z_n_channel, merge_n_imgs, merge_n_imgs)), dtype=ms.float32)\n", + " batch_d_loss_val, batch_g_loss_val = define_evaluation(\n", + " u_r10_validation, u_r5_validation, u_r3_validation, u_r1_validation,\n", + " cp_fluc_validation, Re_c_validation, scaling_input_validation, latent_z_input_validation, idx_)\n", + "\n", + " # predict the d_loss and g_loss in testing\n", + " idx__ = np.random.randint(0, sample_num_test, batch_size)\n", + " latent_z_input_test = Tensor(np.random.normal(0, 1, (\n", + " batch_size, latent_z_n_channel, merge_n_imgs, merge_n_imgs)), dtype=ms.float32)\n", + " batch_d_loss_test, batch_g_loss_test = define_evaluation(\n", + " u_r10_test, u_r5_test, u_r3_test, u_r1_test,\n", + " cp_fluc_test, Re_c_test, scaling_input_test, latent_z_input_test, idx__)\n", + "\n", + " end = time.perf_counter()\n", + " time_consumption = (end - start) // 60 # covert s. to min.\n", + " # Plot the progress and index[0] is the weighted loss\n", + " print(\n", + " \"%d:[%d min] [D:%.2f, %.2f, %.2f, %.2f, %.2f] [G:%.2f, %.2f, %.2f, %.2f, %.2f] \"\n", + " \"[L2_train:%.5f, %.5f, %.5f, %.5f] [L2_val:%.5f, %.5f, %.5f, %.5f] [L2_test:%.5f, %.5f, %.5f, %.5f]\" %\n", + " (epoch, time_consumption,\n", + " ave_d_loss_train[0],\n", + " ave_d_loss_train[1] + ave_d_loss_train[5] + ave_d_loss_train[9],\n", + " ave_d_loss_train[2] + ave_d_loss_train[6] + ave_d_loss_train[10],\n", + " ave_d_loss_train[3] + ave_d_loss_train[7] + ave_d_loss_train[11],\n", + " ave_d_loss_train[4] + ave_d_loss_train[8] + ave_d_loss_train[12],\n", + " ave_g_loss_train[0],\n", + " ave_g_loss_train[1], ave_g_loss_train[2], ave_g_loss_train[3], ave_g_loss_train[4],\n", + " ave_g_loss_train[5], ave_g_loss_train[6], ave_g_loss_train[7], ave_g_loss_train[8],\n", + " batch_g_loss_val[5], batch_g_loss_val[6], batch_g_loss_val[7], batch_g_loss_val[8],\n", + " batch_g_loss_test[5], batch_g_loss_test[6], batch_g_loss_test[7], batch_g_loss_test[8]))\n", + "\n", + " losslog.append(np.concatenate( # 60\n", + " ([ave_d_loss_train[i] for i in range(13)], # 13\n", + " [batch_d_loss_val[i] for i in range(13)], # 13\n", + " [batch_d_loss_test[i] for i in range(13)], # 13\n", + " [ave_g_loss_train[i] for i in range(13)], # 13\n", + " [batch_g_loss_val[i] for i in range(13)], # 13\n", + " [batch_g_loss_test[i] for i in range(13)]))) # 13\n", + "\n", + " # If at save interval => save generated image samples\n", + " if epoch % sample_interval == 0:\n", + " sample_images(epoch)\n", + " ms.save_checkpoint(gan_model, \"training_results/model/gan_%d\" % epoch + \".ckpt\")\n", + " ms.save_checkpoint(critic_model, \"training_results/model/critic_%d\" % epoch + \".ckpt\")\n", + " np.savetxt('training_results/loss/loss.txt', losslog, fmt='%.4f')\n", + " del ave_d_loss_train, ave_g_loss_train, batch_d_loss_val, batch_g_loss_val, batch_d_loss_test, batch_g_loss_test\n", + " gc.collect()" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "pycharm": { + "name": "#%% md\n" + } + }, + "source": [ + "## **结果可视化**\n", + "\n", + "300epoch训练下，网络生成训练集速度U和V方向各尺度速度场的预测值与真实值如下图所示：\n", + "\n", + " \n", + "\n", + "300epoch训练下，网络生成测试集速度U和V方向各尺度速度场的预测值与真实值如下图所示：\n", + "\n", + " \n", + "\n", + "其中一行为一个样本，每张图两列一组共四组，从左往右尺度依次减小，每组左边为生成值右边为真值。\n" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "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.13" + }, + "vscode": { + "interpreter": { + "hash": "b9063439a3781aed32d6b0dd4804a0c8b51ecec7893a0f31b99846bc91ef39eb" + } + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} diff --git a/MindFlow/applications/research/cascade_net/README.md b/MindFlow/applications/research/cascade_net/README.md new file mode 100644 index 0000000000000000000000000000000000000000..2a21e5199ca85d3c1da4909df8fc9a46dc410c99 --- /dev/null +++ b/MindFlow/applications/research/cascade_net/README.md @@ -0,0 +1,81 @@ +# Prediction of spatiotemporal field of pulsation velocity in cylindrical wake by Cascade Net + +## Background + +In the process of turbulent spatiotemporal evolution, the pulsating velocity field includes a series of important fluid physical processes, such as separation, transition, and energy transfer. At high Reynolds numbers, the pulsating velocity field exhibits significant nonlinear characteristics. There are vortex structures in turbulent wake that range from maximum to minimum scales, and these fluid motion patterns constitute complex flow field structural characteristics. The process of energy transfer from large-scale structures to small-scale structures in these flow field structures is called the energy cascade physics principle. Inspired by this principle, the small-scale prediction problem can be transformed into a step-by-step prediction problem from large-scale to small-scale. + +## Model framework + +The model framework is as shown in the following figure: + +![Cascade-Net](images/Cascade-Net.png) + +Where, the generator is a U-Net structure with spatial and channel attention gates, and its framework is shown in the following figure: + +![The U-Net structure of the generator with spatial and channel attention gates](images/The_U-Net_structure_of_the_generator_with_spatial_and_channel_attention_gates.png) + +And for the spatial attention gate *S* and channel attention gate *C*, the integration diagrams are shown as follows: + +![Spatial attention gate S](images/Spatial_attention_gate_S.png) + +![Channel attention gate C](images/Channel_attention_gate_C.png) + +## QuickStart + +Dataset download link: [Cascade_Net/dataset](https://download-mindspore.osinfra.cn/mindscience/mindflow/dataset/applications/research/Cascade_Net/), Save the dataset to `./dataset`. + +The case provides two training methods + +- Run Option 1: Call `train.py` from command line + + ```python + # Call `train.py` from command line + python train.py --config_path ./config/Cascade-Net.yaml + + ``` + + `--config_path` indicates the path of the parameter file. Default path "./config/config.yaml". + + In the "./config/config.yaml" parameter file: + + 'lambda_GP' represents the gradient penalty coefficient, with a default value of 10; + + 'critic_model_lr' represents the learning rate of discriminator, with a default value of 0.00025; + + 'gan_model_lr' represents the learning rate of generator, with a default value of 0.004； + + + +- Run Option 2: Run Jupyter Notebook + + You can use [Chinese](./Cascade-Net_CN.ipynb) or [English](./Cascade-Net.ipynb) Jupyter Notebook to run the training and evaluation code line-by-line. + +## Performance + +| Parameters | GPU | NPU | +|:-------------------------:|:-----------------------:|:------------------:| +| hardware | NVIDIA 3090(memory 24G) | Ascend(memory 32G) | +| MindSpore version | 2.2.14 | 2.2.14 | +| data_size | 10792 | 10792 | +| batch_size | 128 | 128 | +| epochs | 300 | 300 | +| optimizer | RMSProp | RMSProp | +| scale I train loss(MAE) | 3.4e-2 | 4.6e-2 | +| scale I test loss(MAE) | 5.4e-2 | 5.2e-2 | +| scale II train loss(MAE) | 3.2e-2 | 3.6e-2 | +| scale II test loss(MAE) | 5.2e-2 | 5.0e-2 | +| scale III train loss(MAE) | 3.3e-2 | 3.7e-2 | +| scale III test loss(MAE) | 5.4e-2 | 5.2e-2 | +| scale IV train loss(MAE) | 3.6e-2 | 4.1e-2 | +| scale IV test loss(MAE) | 5.5e-2 | 5.4e-2 | +| speed(s/step) | 2.60 | 3.94 | + +## Contributor + +gitee id: [GQEm](https://gitee.com/guoqicheng2024) +email: qicheng@isrc.iscas.ac.cn + +## Reference + +Mi J, Jin X, Li H. Cascade-Net for predicting cylinder wake at Reynolds numbers ranging from subcritical to supercritical regime[J]. Physics of Fluids, 2023, 35: 075132. https://doi.org/10.1063/5.0155649 + diff --git a/MindFlow/applications/research/cascade_net/README_CN.md b/MindFlow/applications/research/cascade_net/README_CN.md new file mode 100644 index 0000000000000000000000000000000000000000..6b815019b25bcec4d2067b84688535c2a96e2d3c --- /dev/null +++ b/MindFlow/applications/research/cascade_net/README_CN.md @@ -0,0 +1,82 @@ +# Cascade-Net预测圆柱尾迹脉动速度时空场 + +## 背景介绍 + +在湍流时空演化过程中，脉动速度场包含了一系列重要的流体物理过程，如分离、转捩和能量传递。在高雷诺数下，脉动速度场表现出明显的非线性特征。湍流尾迹中存在着从最大尺度到最小尺度的涡结构，这些流体运动模式构成了复杂的流场结构特征。在这些流场结构中，能量从大尺度结构向小尺度结构转移的过程被称为能量级串物理原理。受到这一原理的启发，可以将小尺度预测问题转化为由大尺度向小尺度逐级预测问题。 + +## 模型框架 + +模型框架如下图所示： + +![Cascade-Net](images/Cascade-Net.png) + +图中，Generater为具有空间和通道注意力门的U-Net结构，其框架如下图所示： + +![The U-Net structure of the generator with spatial and channel attention gates](images/The_U-Net_structure_of_the_generator_with_spatial_and_channel_attention_gates.png) + +其中空间注意力门*S*和通道注意力门*C*示意图如下： + +![Spatial attention gate S](images/Spatial_attention_gate_S.png) + +![Channel attention gate C](images/Channel_attention_gate_C.png) + +## 快速开始 + +数据集下载地址：[Cascade_Net/dataset](https://download-mindspore.osinfra.cn/mindscience/mindflow/dataset/applications/research/Cascade_Net/). 将数据集保存在`./dataset`路径下。 + +案例提供两种训练方式 + +- 训练方式一：在命令行中调用`train.py`脚本 + + ```python + # 在命令行调用train.py进行训练示例 + python train.py --config_path ./config/Cascade-Net.yaml + + ``` + + `--config_path`表示配置文件的路径，默认值为 "./config/config.yaml"。 + + 在 "./config/config.yaml" 配置文件中： + + 'lambda_GP' 表示梯度惩罚系数，默认值为10； + + 'critic_model_lr' 表示判别器学习率，默认值为0.00025； + + 'gan_model_lr' 表示生成器学习率，默认值为0.004； + + + +- 训练方式二：运行Jupyter Notebook + + 您可以使用[中文版](./Cascade-Net_CN.ipynb) 和[英文版](./Cascade-Net.ipynb) Jupyter Notebook逐行运行训练和验证代码。 + +## 性能 + +| 参数 | GPU | NPU | +|:-------------------:|:-----------------------:|:------------------:| +| 硬件 | NVIDIA 3090(memory 24G) | Ascend(memory 32G) | +| MindSpore版本 | 2.2.14 | 2.2.14 | +| 数据大小 | 10792 | 10792 | +| batch大小 | 128 | 128 | +| 训练步数 | 300 | 300 | +| 优化器 | RMSProp | RMSProp | +| scale I 训练精度(MAE) | 3.4e-2 | 4.6e-2 | +| scale I 测试精度(MAE) | 5.4e-2 | 5.2e-2 | +| scale II 训练精度(MAE) | 3.2e-2 | 3.6e-2 | +| scale II 测试精度(MAE) | 5.2e-2 | 5.0e-2 | +| scale III 训练精度(MAE) | 3.3e-2 | 3.7e-2 | +| scale III 测试精度(MAE) | 5.4e-2 | 5.2e-2 | +| scale IV 训练精度(MAE) | 3.6e-2 | 4.1e-2 | +| scale IV 测试精度(MAE) | 5.5e-2 | 5.4e-2 | +| 性能(s/step) | 2.60 | 3.94 | + + + +## 贡献者 + +gitee id: [GQEm](https://gitee.com/guoqicheng2024) +email: qicheng@isrc.iscas.ac.cn + +## 参考文献 + +Mi J, Jin X, Li H. Cascade-Net for predicting cylinder wake at Reynolds numbers ranging from subcritical to supercritical regime[J]. Physics of Fluids, 2023, 35: 075132. https://doi.org/10.1063/5.0155649 \ No newline at end of file diff --git a/MindFlow/applications/research/cascade_net/config/Cascade-Net.yaml b/MindFlow/applications/research/cascade_net/config/Cascade-Net.yaml new file mode 100644 index 0000000000000000000000000000000000000000..a69c8760d4a265580305d7ac9f2fef88a39aa49c --- /dev/null +++ b/MindFlow/applications/research/cascade_net/config/Cascade-Net.yaml @@ -0,0 +1,28 @@ +#data: +batch_size: 128 +n_channel_p: 3 # Cp = (None, 128, 128, 3) +n_channel_u: 2 # u = (None, 128, 128, 2) including u and v +latent_z_n_channel: 50 # z = (N,8,8,50) + +#model: +#critic: +lambda_GP: 10 +critic_model_lr: 0.00025 +n_critic: 3 + +#ganerater: +lambda_L2_u: 1000 +lambda_L2_gradu: 100 +gan_model_lr: 0.004 + +#summary: +plot_n: 5 +n_imgs: 128 +merge_n_imgs: 8 +epochs: 60000 +sample_interval: 10 + + + + + diff --git a/MindFlow/applications/research/cascade_net/dataset/.keep b/MindFlow/applications/research/cascade_net/dataset/.keep new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/MindFlow/applications/research/cascade_net/images/Cascade-Net.png b/MindFlow/applications/research/cascade_net/images/Cascade-Net.png new file mode 100644 index 0000000000000000000000000000000000000000..352770d27263c19085ea20561c696802a4f9364b Binary files /dev/null and b/MindFlow/applications/research/cascade_net/images/Cascade-Net.png differ diff --git a/MindFlow/applications/research/cascade_net/images/Channel_attention_gate_C.png b/MindFlow/applications/research/cascade_net/images/Channel_attention_gate_C.png new file mode 100644 index 0000000000000000000000000000000000000000..989b919428de75419c22b062616e8bb4b5ed50e7 Binary files /dev/null and b/MindFlow/applications/research/cascade_net/images/Channel_attention_gate_C.png differ diff --git a/MindFlow/applications/research/cascade_net/images/Sample_cases_for_training_validation_and_testing.png b/MindFlow/applications/research/cascade_net/images/Sample_cases_for_training_validation_and_testing.png new file mode 100644 index 0000000000000000000000000000000000000000..8ddb0b7030ee1c4150f79969182863547a849cb6 Binary files /dev/null and b/MindFlow/applications/research/cascade_net/images/Sample_cases_for_training_validation_and_testing.png differ diff --git a/MindFlow/applications/research/cascade_net/images/Spatial_attention_gate_S.png b/MindFlow/applications/research/cascade_net/images/Spatial_attention_gate_S.png new file mode 100644 index 0000000000000000000000000000000000000000..df063463e2c3727a400eebbb55cccb9e05f191ea Binary files /dev/null and b/MindFlow/applications/research/cascade_net/images/Spatial_attention_gate_S.png differ diff --git a/MindFlow/applications/research/cascade_net/images/The_U-Net_structure_of_the_generator_with_spatial_and_channel_attention_gates.png b/MindFlow/applications/research/cascade_net/images/The_U-Net_structure_of_the_generator_with_spatial_and_channel_attention_gates.png new file mode 100644 index 0000000000000000000000000000000000000000..d174b3d1e20abae34bf05227f5592b94ee4b5bbc Binary files /dev/null and b/MindFlow/applications/research/cascade_net/images/The_U-Net_structure_of_the_generator_with_spatial_and_channel_attention_gates.png differ diff --git a/MindFlow/applications/research/cascade_net/images/test_u.png b/MindFlow/applications/research/cascade_net/images/test_u.png new file mode 100644 index 0000000000000000000000000000000000000000..b988c00d9a0a6388646e25bafc5243e3948d90a2 Binary files /dev/null and b/MindFlow/applications/research/cascade_net/images/test_u.png differ diff --git a/MindFlow/applications/research/cascade_net/images/test_v.png b/MindFlow/applications/research/cascade_net/images/test_v.png new file mode 100644 index 0000000000000000000000000000000000000000..20b145b9bb8260f328920cc7c16698846802e17f Binary files /dev/null and b/MindFlow/applications/research/cascade_net/images/test_v.png differ diff --git a/MindFlow/applications/research/cascade_net/images/train_u.png b/MindFlow/applications/research/cascade_net/images/train_u.png new file mode 100644 index 0000000000000000000000000000000000000000..c40f4abad91788f70658c2d59f0145f84ad96da9 Binary files /dev/null and b/MindFlow/applications/research/cascade_net/images/train_u.png differ diff --git a/MindFlow/applications/research/cascade_net/images/train_v.png b/MindFlow/applications/research/cascade_net/images/train_v.png new file mode 100644 index 0000000000000000000000000000000000000000..29f7e442302d4d27946d6350046530176e15aa4b Binary files /dev/null and b/MindFlow/applications/research/cascade_net/images/train_v.png differ diff --git a/MindFlow/applications/research/cascade_net/src/Compound_Model.py b/MindFlow/applications/research/cascade_net/src/Compound_Model.py new file mode 100644 index 0000000000000000000000000000000000000000..bfc4912e84d79e635a707bef01dcaa29cfd83df5 --- /dev/null +++ b/MindFlow/applications/research/cascade_net/src/Compound_Model.py @@ -0,0 +1,204 @@ +# Copyright 2023 Huawei Technologies Co., Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +import mindspore.nn as nn +from mindspore import ops +from .ResUnetWithAttentionSpatialChannelGenerator import AttentionResUnet +from .MultiscaleDiscriminator import MultiscaleDiscriminator +from functools import partial +from .Loss import GradientPenaltyLoss + + +class DefineMerge(nn.Cell): + def __init__(self, merge_n_imgs, merge_n_channel): + super(DefineMerge, self).__init__() + self.merge_n_imgs = merge_n_imgs + self.merge_n_channel = merge_n_channel + self.define_merge_SC = nn.SequentialCell( + nn.Dense(21, self.merge_n_imgs * self.merge_n_imgs * self.merge_n_channel), + nn.BatchNorm1d(self.merge_n_imgs * self.merge_n_imgs * self.merge_n_channel), + nn.ReLU() + ) + self.Concat = ops.Concat(axis=1) + + def construct(self, in_u_scaling, in_Re): # dim of in_merge_map is (batch_size, 8, 8, 1) + # in_u_scaling # Inputs scaling u, [,20] + # in_Re # Inputs scaling u, [,1] + # in_merge # mege inputs, [,11] + in_merge = self.Concat((in_u_scaling, in_Re)) + in_merge_map = self.define_merge_SC(in_merge) # [,11] -> [,8*8*1] = [,64] + in_merge_map = ops.reshape(in_merge_map, + (in_merge_map.shape[0], self.merge_n_channel, self.merge_n_imgs, + self.merge_n_imgs)) # [,64] -> [, 8, 8, 1] + return in_merge_map + + +def sub_model_block_relu(n_channel_u, in_dim, filter_num_g, filter_num_d, filter_size, g_merge, g_latent_z, d_multi, + d_single): + sub_g_model = AttentionResUnet( + input_channel=in_dim, + output_mask_channel=n_channel_u, + filter_num=filter_num_g, + filter_size=filter_size, + merge=g_merge, latent_z=g_latent_z) + sub_d_model = MultiscaleDiscriminator( + input_channel=n_channel_u, + filter_num=filter_num_d, + multi_scale=d_multi, + single_scale=d_single) + return sub_g_model, sub_d_model + + +def init_sub_model(n_channel_p, n_channel_u): + merge_model = DefineMerge(8, 1) + g_model_I, d_model_I = sub_model_block_relu( + n_channel_u=n_channel_u, + in_dim=n_channel_p, + filter_num_g=8, + filter_num_d=16, + filter_size=3, + g_merge=True, + g_latent_z=False, + d_multi=False, + d_single=True + ) + g_model_II, d_model_II = sub_model_block_relu( + n_channel_u=n_channel_u, + in_dim=(n_channel_p + n_channel_u), + filter_num_g=8, + filter_num_d=8, + filter_size=3, + g_merge=True, + g_latent_z=False, + d_multi=True, + d_single=False + ) + g_model_III, d_model_III = sub_model_block_relu( + n_channel_u=n_channel_u, + in_dim=(n_channel_p + n_channel_u + n_channel_u), + # 2022.11.19改动: gen_u_III from gen_u_I + gen_u_II + filter_num_g=8, + filter_num_d=8, + filter_size=3, + g_merge=True, + g_latent_z=False, + d_multi=True, + d_single=False + ) + g_model_IV, d_model_IV = sub_model_block_relu( + n_channel_u=n_channel_u, + in_dim=(n_channel_p + n_channel_u + n_channel_u + n_channel_u), + # 2022.11.19改动: gen_u_III from gen_u_I + gen_u_II + filter_num_g=8, + filter_num_d=8, + filter_size=2, + g_merge=False, + g_latent_z=True, + d_multi=True, + d_single=False + ) + return merge_model, g_model_I, d_model_I, g_model_II, d_model_II, g_model_III, d_model_III, g_model_IV, d_model_IV + + +class DefineCompoundCritic(nn.Cell): + def __init__(self, n_channel_p, n_channel_u, batch_size, d_model_I, d_model_II, d_model_III, d_model_IV): + super(DefineCompoundCritic, self).__init__() + self.n_channel_p = n_channel_p + self.n_channel_u = n_channel_u + self.batch_size = batch_size + + self.d_model_I = d_model_I + self.d_model_II = d_model_II + self.d_model_III = d_model_III + self.d_model_IV = d_model_IV + + self.gradient_penalty_loss = GradientPenaltyLoss() + self.Concat = ops.Concat(axis=1) + + def construct(self, input): + in_u_I = input[0] # ,128,128,2 + in_u_II = input[1] # ,128,128,2 + in_u_III = input[2] # ,128,128,2 + in_u_IV = input[3] # ,128,128,2 + in_merge = input[4] + gen_u_I = input[5] + gen_u_II = input[6] + gen_u_III = input[7] + gen_u_IV = input[8] + + alpha = ops.rand((self.batch_size, 2, 1, 1)) # uv + + d_on_real_I = self.d_model_I(in_u_I, in_merge) + d_on_fake_I = self.d_model_I(gen_u_I, in_merge) + diff_fv = in_u_I - gen_u_I + interpolated_img_I = gen_u_I + alpha * diff_fv + d_on_interp_I = self.d_model_I(interpolated_img_I, in_merge) + partial_gp_loss_I = partial(self.gradient_penalty_loss, d_model=self.d_model_I, interpolated_img=interpolated_img_I, in_merge=in_merge) + + d_on_real_II = self.d_model_II(in_u_II, in_merge) + d_on_fake_II = self.d_model_II(gen_u_II, in_merge) + diff_fv = in_u_II - gen_u_II + interpolated_img_II = gen_u_II + alpha * diff_fv + d_on_interp_II = self.d_model_II(interpolated_img_II, in_merge) + partial_gp_loss_II = partial(self.gradient_penalty_loss, d_model=self.d_model_II, interpolated_img=interpolated_img_II, in_merge=in_merge) + + d_on_real_III = self.d_model_III(in_u_III, in_merge) + d_on_fake_III = self.d_model_III(gen_u_III, in_merge) + diff_fv = in_u_III - gen_u_III + interpolated_img_III = gen_u_III + alpha * diff_fv + d_on_interp_III = self.d_model_III(interpolated_img_III, in_merge) + partial_gp_loss_III = partial(self.gradient_penalty_loss, d_model=self.d_model_III, interpolated_img=interpolated_img_III, in_merge=in_merge) + + d_on_real_IV = self.d_model_IV(in_u_IV, in_merge) + d_on_fake_IV = self.d_model_IV(gen_u_IV, in_merge) + diff_fv = in_u_IV - gen_u_IV + interpolated_img_IV = gen_u_IV + alpha * diff_fv + d_on_interp_IV = self.d_model_IV(interpolated_img_IV, in_merge) + partial_gp_loss_IV = partial(self.gradient_penalty_loss, d_model=self.d_model_IV, interpolated_img=interpolated_img_IV, in_merge=in_merge) + + return (d_on_real_I, d_on_real_II, d_on_real_III, d_on_real_IV, + d_on_fake_I, d_on_fake_II, d_on_fake_III, d_on_fake_IV, + d_on_interp_I, d_on_interp_II, d_on_interp_III, d_on_interp_IV, + partial_gp_loss_I, partial_gp_loss_II, partial_gp_loss_III, partial_gp_loss_IV) + + +class DefineCompoundGan(nn.Cell): + def __init__(self, n_channel_p, n_channel_v, merge_model, g_model_I, g_model_II, g_model_III, g_model_IV): + super(DefineCompoundGan, self).__init__() + self.n_channel_p = n_channel_p + self.n_channel_u = n_channel_v + + self.merge_model = merge_model + self.g_model_I = g_model_I + self.g_model_II = g_model_II + self.g_model_III = g_model_III + self.g_model_IV = g_model_IV + + self.Concat = ops.Concat(axis=1) + + def construct(self, input): + in_cp = input[0] # 3 + in_Re = input[1] # Conditions Re, [,1] + in_u_scaling = input[2] # Inputs scaling u, [,20] + in_z = input[3] # latent_z inputs, [8,8,50] + in_merge = self.merge_model(in_Re, in_u_scaling) # merge inputs, [8,8,1] + + # connect the input and output of the generator + gen_u_I = self.g_model_I(in_cp, in_merge) + gen_u_II = self.g_model_II(self.Concat((in_cp, gen_u_I)), in_merge) + gen_u_III = self.g_model_III(self.Concat((in_cp, gen_u_I, gen_u_II)), in_merge) + gen_u_IV = self.g_model_IV(self.Concat((in_cp, gen_u_I, gen_u_II, gen_u_III)), in_merge, in_z) + + return (gen_u_I, gen_u_II, gen_u_III, gen_u_IV, in_merge, + gen_u_I, gen_u_II, gen_u_III, gen_u_IV, in_merge) diff --git a/MindFlow/applications/research/cascade_net/src/Loss.py b/MindFlow/applications/research/cascade_net/src/Loss.py new file mode 100644 index 0000000000000000000000000000000000000000..46e685abf30147b8d0a639008ca7b55567648719 --- /dev/null +++ b/MindFlow/applications/research/cascade_net/src/Loss.py @@ -0,0 +1,94 @@ +# Copyright 2023 Huawei Technologies Co., Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +from mindspore import ops +import mindspore.nn as nn + + +class WassersteinLoss(nn.Cell): + def __init__(self): + super(WassersteinLoss, self).__init__() + + def construct(self, y_pred, y_true): + loss = y_true * y_pred + loss_mean = ops.mean(loss, axis=(1, 2, 3, 4)) + return loss_mean + + +class GradLoss(nn.Cell): + def __init__(self): + super(GradLoss, self).__init__() + self.dx = 0.0242 # updated + self.dy = 0.0291 + self.Concat = ops.Concat(axis=2) + + def CalGrad(self, u_mat): + ''' + :param u_mat: the input, [batch_size,128,128,2] including u and v + ''' + group_00 = u_mat[:, :, :, 0:-2, 0:-2] # [batch_size,2,126,126] + group_01 = u_mat[:, :, :, 0:-2, 1:-1] + group_02 = u_mat[:, :, :, 0:-2, 2:] + group_10 = u_mat[:, :, :, 1:-1, 0:-2] + group_11 = u_mat[:, :, :, 1:-1, 1:-1] + group_12 = u_mat[:, :, :, 1:-1, 2:] + group_20 = u_mat[:, :, :, 2:, 0:-2] + group_21 = u_mat[:, :, :, 2:, 1:-1] + group_22 = u_mat[:, :, :, 2:, 2:] + + grad_x = (group_20 + 2 * group_21 + group_22 - group_00 - 2 * group_01 - group_02) / ( + 4 * 2 * self.dx) # [batch_size,126,126,2] du_dx,dv_dx + grad_y = (group_02 + 2 * group_12 + group_22 - group_00 - 2 * group_10 - group_20) / ( + 4 * 2 * self.dy) # [batch_size,126,126,2] du_dy,dv_dy + grad_all = self.Concat((grad_x, grad_y)) + return grad_all + + def construct(self, y_pred, y_true): + grad_y_true = self.CalGrad(y_true) + grad_y_fake = self.CalGrad(y_pred) + gradloss = ops.mean(ops.square(grad_y_true - grad_y_fake), axis=(1, 2, 3, 4)) + return gradloss + + +class GradientPenaltyLoss(nn.Cell): + """ + Computes gradient penalty based on prediction and weighted real / fake samples + """ + def __init__(self): + super(GradientPenaltyLoss, self).__init__() + self.ReduceSum = ops.ReduceSum(keep_dims=False) + self.Square = ops.Square() + self.Sqrt = ops.Sqrt() + self.grad_op = ops.GradOperation() + + def construct(self, d_model, interpolated_img, in_merge): + # first get the gradients: + # assuming: - that y_pred has dimensions (batch_size, patch_size) + # - averaged_samples has dimensions (batch_size, nbr_features) + # gradients afterwards has dimension (batch_size, nbr_features), basically + # a list of nbr_features-dimensional gradient vectors + gradient_function = self.grad_op(d_model) + gradients = gradient_function(interpolated_img, in_merge)[0] + # gradients = ms.grad(d_model, 0)(interpolated_img, in_merge) + # compute the euclidean norm by squaring ... + gradients_sqr = self.Square(gradients) + # ... summing over the rows ... + gradients_sqr_sum = self.ReduceSum(gradients_sqr, axis=tuple(range(2, len(gradients_sqr.shape)))) + # ... and sqrt + gradient_l2_norm = self.Sqrt(gradients_sqr_sum) + # compute lambda * (1 - ||grad||)^2 still for each single sample + gradient_penalty = self.Square(1 - gradient_l2_norm) + # return the mean as loss over all the batch samples + + return ops.mean(gradient_penalty) \ No newline at end of file diff --git a/MindFlow/applications/research/cascade_net/src/MultiscaleDiscriminator.py b/MindFlow/applications/research/cascade_net/src/MultiscaleDiscriminator.py new file mode 100644 index 0000000000000000000000000000000000000000..74f0db5439d5e23e78a1eea071d0f7480bed30b0 --- /dev/null +++ b/MindFlow/applications/research/cascade_net/src/MultiscaleDiscriminator.py @@ -0,0 +1,104 @@ +# Copyright 2023 Huawei Technologies Co., Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +import mindspore.nn as nn +from mindspore import ops, Tensor + + +class AveragePool2d(nn.Cell): + def __init__(self, kernel_size=(20, 20), stride=1, pad_mode='same'): + super(AveragePool2d, self).__init__() + self.kernel_size = kernel_size + self.stride = stride + self.pad_mode = pad_mode + weight = ops.zeros((2, 2, self.kernel_size[0], self.kernel_size[1])) + weight[0, 0, :, :] = 1 + weight[1, 1, :, :] = 1 + self.weight = weight / self.kernel_size[0] / self.kernel_size[1] + + def construct(self, x): + out = ops.conv2d(x, weight=self.weight, stride=self.stride, pad_mode=self.pad_mode) + return out + + +class DefineDiscriminatorBlock(nn.Cell): + def __init__(self, in_channel, filter_num): # filter_num 16 for best + super(DefineDiscriminatorBlock, self).__init__() + self.discriminator_sequential = nn.SequentialCell( + # C8 + nn.Conv2d(in_channel, filter_num, (3, 3), stride=(2, 2), has_bias=True), # (None, 64, 64, 8) + nn.ReLU(), + nn.Dropout(p=0.25), # 缺少training=True参数 + # C16 + nn.Conv2d(filter_num, 2 * filter_num, (3, 3), stride=(2, 2), has_bias=True), # (None, 32, 32, 16) + nn.ReLU(), + nn.Dropout(p=0.25), + # C32 + nn.Conv2d(2 * filter_num, 4 * filter_num, (3, 3), stride=(2, 2), has_bias=True), # (None, 16, 16, 32) + nn.ReLU(), + nn.Dropout(p=0.25), + # C64 + nn.Conv2d(4 * filter_num, 8 * filter_num, (3, 3), stride=(2, 2), has_bias=True), # (None, 8, 8, 64) + nn.ReLU(), + nn.Dropout(p=0.25), + # second last output layer + nn.Conv2d(8 * filter_num, 16 * filter_num, (3, 3), has_bias=True), # (None, 8, 8, 128) + nn.ReLU(), + ) + self.Conv = nn.Conv2d(16 * filter_num + 1, 1, (3, 3), has_bias=True) + self.Concat = ops.Concat(axis=1) + + def construct(self, in_D, in_merge): + # in_D = input_data[0] # (None,128,128,2) uv + # in_merge = input_data[1] # [,8,8,1] + # source image input + d = self.discriminator_sequential(in_D) + d = self.Concat((d, in_merge)) + # patch output + patch_out = self.Conv(d) # (None, 8, 8, 1) + return patch_out + + +class MultiscaleDiscriminator(nn.Cell): + def __init__(self, input_channel, filter_num, multi_scale=False, single_scale=False): # filter_num 16 for best + super(MultiscaleDiscriminator, self).__init__() + self.D_weight = Tensor([1, 5, 10]) # ,20,10,and raw + self.single_scale = single_scale + self.multi_scale = multi_scale + self.Discriminator1 = DefineDiscriminatorBlock(input_channel, filter_num) + if self.multi_scale: + self.AvePool20 = AveragePool2d(kernel_size=(20, 20), stride=1, pad_mode='same') + self.AvePool10 = AveragePool2d(kernel_size=(10, 10), stride=1, pad_mode='same') + # self.AvePool20 = nn.AvgPool2d(kernel_size=(20, 20), stride=1, pad_mode='same') + # self.AvePool10 = nn.AvgPool2d(kernel_size=(10, 10), stride=1, pad_mode='same') + self.Discriminator2 = DefineDiscriminatorBlock(input_channel, filter_num) + self.Discriminator3 = DefineDiscriminatorBlock(input_channel, filter_num) + self.Concat = ops.Concat(axis=1) + + def construct(self, in_D, in_merge): + # in_D = input_data[0] # (None,128,128,2) uv + # in_merge = input_data[1] # [,8,8,1] + if self.single_scale: + out_final = self.Discriminator1(in_D, in_merge) # (None, 8,8,1) + if self.multi_scale: + ur20 = self.AvePool20(in_D) + out_map_20 = self.Discriminator1(ur20, in_merge) # (None, 8,8,1) + ur10 = self.AvePool10(in_D) + out_map_10 = self.Discriminator2(ur10, in_merge) # (None, 8,8,1) + out_map_raw = self.Discriminator3(in_D, in_merge) # (None, 8,8,1) + out_map_concat = self.Concat((self.D_weight[0] * out_map_20, + self.D_weight[1] * out_map_10, + self.D_weight[2] * out_map_raw)) + out_final = ops.mean(out_map_concat, axis=1, keep_dims=True) + return out_final diff --git a/MindFlow/applications/research/cascade_net/src/ResUnetWithAttentionSpatialChannelGenerator.py b/MindFlow/applications/research/cascade_net/src/ResUnetWithAttentionSpatialChannelGenerator.py new file mode 100644 index 0000000000000000000000000000000000000000..23626ac9ea13d9f1ca2b4a6b638c2c19315db9fe --- /dev/null +++ b/MindFlow/applications/research/cascade_net/src/ResUnetWithAttentionSpatialChannelGenerator.py @@ -0,0 +1,295 @@ +# Copyright 2023 Huawei Technologies Co., Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +import mindspore.nn as nn +from mindspore import ops + + +class DoubleConvLayer(nn.Cell): + """ + construction of a double convolutional layer using + SAME padding + RELU nonlinear activation function + :param x: input + :param filter_size: size of convolutional filter + :param size: number of filters + :param dropout: FLAG & RATE of dropout. + if < 0 dropout cancelled, if > 0 set as the rate + :param batch_norm: flag of if batch_norm used, + if True batch normalization + :return: output of a double convolutional layer + """ + def __init__(self, in_channels, filter_size, size, dropout): # batch_norm all T + super(DoubleConvLayer, self).__init__() + self.dropout = dropout + self.double_conv_sequential = nn.SequentialCell( + nn.Conv2d(in_channels, size, (filter_size, filter_size), has_bias=True), + nn.BatchNorm2d(size, eps=1e-3, momentum=0.99), + nn.ReLU(), + nn.Conv2d(size, size, (filter_size, filter_size), has_bias=True), + nn.BatchNorm2d(size, eps=1e-3, momentum=0.99), + nn.ReLU(), + ) + self.Conv_size1 = nn.Conv2d(in_channels, size, (1, 1), has_bias=True) + self.BatchNormal = nn.BatchNorm2d(size, eps=1e-3, momentum=0.99) + if dropout > 0: + self.Drop = nn.Dropout(p=dropout) + + def construct(self, x): + conv = self.double_conv_sequential(x) + if self.dropout > 0: + conv = self.Drop(conv) + shortcut = self.Conv_size1(x) + shortcut = self.BatchNormal(shortcut) + res_path = shortcut + conv + return res_path + + +class SeBlock(nn.Cell): + """ + self attention squeeze-excitation block, attention mechanism on channel dimension + :param x: input feature map + :return: attention weighted on channel dimension feature map + """ + def __init__(self, in_channels, out_dim, ratio): + super(SeBlock, self).__init__() + self.out_dim = out_dim + self.SE_sequentialcell = nn.SequentialCell( + nn.BatchNorm1d(in_channels, eps=1e-3, momentum=0.99), + nn.Dense(in_channels, out_dim // ratio, has_bias=True), + nn.ReLU(), + nn.BatchNorm1d(out_dim // ratio, eps=1e-3, momentum=0.99), + nn.Dense(out_dim // ratio, out_dim, has_bias=True), + nn.Sigmoid() + ) + + def construct(self, x): + # Squeeze: global average pooling + x_s = ops.mean(x, axis=(2, 3), keep_dims=False) # N,C + # Excitation: bottom-up top-down FCs + x_e = self.SE_sequentialcell(x_s) + x_e = ops.reshape(x_e, (x_e.shape[0], self.out_dim, 1, 1)) + result = ops.mul(x, x_e) + return result + + +class GatingSignal(nn.Cell): + """ + resize the down layer feature map into the same dimension as the up layer feature map + using 1x1 conv + :param input: down-dim feature map + :param in_channels: in channel number + :param out_size:output channel number + :return: the gating feature map with the same dimension of the up layer feature map + """ + def __init__(self, in_channels, out_size): + super(GatingSignal, self).__init__() + self.gating_signal_sequentialcell = nn.SequentialCell( + nn.Conv2d(in_channels, out_size, (1, 1), has_bias=True), + nn.BatchNorm2d(out_size, eps=1e-3, momentum=0.99), + nn.ReLU() + ) + + def construct(self, x): + x = self.gating_signal_sequentialcell(x) + return x + + +class AttentionBlock(nn.Cell): + """ + self gated attention, attention mechanism on spatial dimension + :param x: input feature map + :param gating: gate signal, feature map from the lower layer + :param inter_shape: intermedium channle numer + :param name: name of attention layer, for output + :return: attention weighted on spatial dimension feature map + """ + def __init__(self, in_channels, inter_shape): + super(AttentionBlock, self).__init__() + conv_imgs = 1024//in_channels + gating_imgs = 512//in_channels + self.Conv_nobias = nn.Conv2d(in_channels, inter_shape, (1, 1)) + self.Conv = nn.Conv2d(in_channels, inter_shape, (1, 1), has_bias=True) + self.Conv2DT = nn.Conv2dTranspose(inter_shape, inter_shape, (3, 3), + stride=(conv_imgs//gating_imgs, conv_imgs//gating_imgs), has_bias=True) # 16//8 + self.ReLU = nn.ReLU() + self.Conv_num1 = nn.Conv2d(inter_shape, 1, (1, 1), has_bias=True) + self.Sigmoid = nn.Sigmoid() + self.Conv_result = nn.Conv2d(in_channels, in_channels, (1, 1), has_bias=True) + self.BatchNormal = nn.BatchNorm2d(in_channels, eps=1e-3, momentum=0.99) + + def construct(self, conv, gating): + shape_x = conv.shape + theta_x = self.Conv_nobias(conv) + phi_g = self.Conv(gating) + upsample_g = self.Conv2DT(phi_g) # 64, 16, 16 + concat_xg = upsample_g + theta_x # 64, 16, 16 + act_xg = self.ReLU(concat_xg) + psi = self.Conv_num1(act_xg) + sigmoid_xg = self.Sigmoid(psi) + upsample_psi = ops.tile(sigmoid_xg, (1, shape_x[1], 1, 1)) + y = ops.mul(upsample_psi, conv) + result = self.Conv_result(y) + result_bn = self.BatchNormal(result) + return result_bn + + +class AttentionResUnet(nn.Cell): + """ + Rsidual UNet construction, with attention gate + convolution: 3*3 SAME padding + pooling: 2*2 VALID padding + upsampling: 3*3 VALID padding + final convolution: 1*1 + :param dropout_rate: FLAG & RATE of dropout. + if < 0 dropout cancelled, if > 0 set as the rate + :param batch_norm: flag of if batch_norm used, + if True batch normalization + :return: model + """ + def __init__(self, input_channel, output_mask_channel, filter_num, filter_size, merge=False, latent_z=False): + super(AttentionResUnet, self).__init__() + self.axis = 1 + self.SE_RATIO = 16 + self.FILTER_SIZE = filter_size + self.FILTER_NUM = filter_num + self.merge = merge + self.latent_z = latent_z + + self.double_conv_layer1 = DoubleConvLayer(input_channel, self.FILTER_SIZE, self.FILTER_NUM, dropout=0.5) + self.MaxPooling1 = nn.MaxPool2d(2, 2) + self.double_conv_layer2 = DoubleConvLayer(self.FILTER_NUM, self.FILTER_SIZE, 2 * self.FILTER_NUM, dropout=0.5) + self.MaxPooling2 = nn.MaxPool2d(2, 2) + self.double_conv_layer3 = DoubleConvLayer(2 * self.FILTER_NUM, self.FILTER_SIZE, 4 * self.FILTER_NUM, dropout=0.5) + self.MaxPooling3 = nn.MaxPool2d(2, 2) + self.double_conv_layer4 = DoubleConvLayer(4 * self.FILTER_NUM, self.FILTER_SIZE, 8 * self.FILTER_NUM, dropout=0.5) + self.MaxPooling4 = nn.MaxPool2d(2, 2) + self.double_conv_layer5 = DoubleConvLayer(8 * self.FILTER_NUM, self.FILTER_SIZE, 16 * self.FILTER_NUM, dropout=0.5) + + if self.latent_z: + self.gating_signal1 = GatingSignal(16 * self.FILTER_NUM + 51, 8 * self.FILTER_NUM) + self.Conv2DTranspose1 = nn.Conv2dTranspose(16 * self.FILTER_NUM + 51, 16 * self.FILTER_NUM, (2, 2), + stride=(2, 2), has_bias=True) + else: + self.gating_signal1 = GatingSignal(16 * self.FILTER_NUM + 1, 8 * self.FILTER_NUM) + self.Conv2DTranspose1 = nn.Conv2dTranspose(16 * self.FILTER_NUM + 1, 16 * self.FILTER_NUM, (2, 2), + stride=(2, 2), has_bias=True) + + self.gating_signal2 = GatingSignal(8 * self.FILTER_NUM, 4 * self.FILTER_NUM) + self.gating_signal3 = GatingSignal(4 * self.FILTER_NUM, 2 * self.FILTER_NUM) + self.gating_signal4 = GatingSignal(2 * self.FILTER_NUM, self.FILTER_NUM) + + self.Conv2DTranspose2 = nn.Conv2dTranspose(8 * self.FILTER_NUM, 8 * self.FILTER_NUM, (2, 2), + stride=(2, 2), has_bias=True) + self.Conv2DTranspose3 = nn.Conv2dTranspose(4 * self.FILTER_NUM, 4 * self.FILTER_NUM, (2, 2), + stride=(2, 2), has_bias=True) + self.Conv2DTranspose4 = nn.Conv2dTranspose(2 * self.FILTER_NUM, 2 * self.FILTER_NUM, (2, 2), + stride=(2, 2), has_bias=True) + + self.attention_block1 = AttentionBlock(8 * self.FILTER_NUM, 8 * self.FILTER_NUM) + self.attention_block2 = AttentionBlock(4 * self.FILTER_NUM, 4 * self.FILTER_NUM) + self.attention_block3 = AttentionBlock(2 * self.FILTER_NUM, 2 * self.FILTER_NUM) + self.attention_block4 = AttentionBlock(self.FILTER_NUM, self.FILTER_NUM) + + self.se_conv_16_layer = SeBlock(16 * self.FILTER_NUM + 8 * self.FILTER_NUM, 16 * self.FILTER_NUM + 8 * self.FILTER_NUM, ratio=self.SE_RATIO) + self.se_conv_32_layer = SeBlock(8 * self.FILTER_NUM + 4 * self.FILTER_NUM, 8 * self.FILTER_NUM + 4 * self.FILTER_NUM, ratio=self.SE_RATIO) + self.se_conv_64_layer = SeBlock(4 * self.FILTER_NUM + 2 * self.FILTER_NUM, 4 * self.FILTER_NUM + 2 * self.FILTER_NUM, ratio=self.SE_RATIO) + self.se_conv_128_layer = SeBlock(2 * self.FILTER_NUM + self.FILTER_NUM, 2 * self.FILTER_NUM + self.FILTER_NUM, ratio=self.SE_RATIO) + + self.up_conv_16_layer = DoubleConvLayer(16 * self.FILTER_NUM + 8 * self.FILTER_NUM, self.FILTER_SIZE, 8 * self.FILTER_NUM, dropout=0.5) + self.up_conv_32_layer = DoubleConvLayer(8 * self.FILTER_NUM + 4 * self.FILTER_NUM, self.FILTER_SIZE, 4 * self.FILTER_NUM, dropout=0.5) + self.up_conv_64_layer = DoubleConvLayer(4 * self.FILTER_NUM + 2 * self.FILTER_NUM, self.FILTER_SIZE, 2 * self.FILTER_NUM, dropout=False) + self.up_conv_128_layer = DoubleConvLayer(2 * self.FILTER_NUM + self.FILTER_NUM, self.FILTER_SIZE, self.FILTER_NUM, dropout=False) + + self.conv_final_layer = nn.Conv2d(self.FILTER_NUM, output_mask_channel, kernel_size=(1, 1), pad_mode='valid', has_bias=True) + self.BatchNormal = nn.BatchNorm2d(output_mask_channel, eps=1e-3, momentum=0.99) + self.Tanh = nn.Tanh() + self.Concat = ops.Concat(axis=self.axis) + + def construct(self, inputs, merge_inputs=None, z_inputs=None): + # input data + # dimension of the image depth + # inputs # (,3,128,128) + + # Down-sampling layers + # DownRes 1, double residual convolution + pooling dropout, batch_norm=False + conv_128 = self.double_conv_layer1(inputs) # (8,128,128) # comments are FILTER_NUM=8 + pool_64 = self.MaxPooling1(conv_128) # (8,64,64) + # DownRes 2 + conv_64 = self.double_conv_layer2(pool_64) # (16,64,64) + pool_32 = self.MaxPooling2(conv_64) # (16,32,32) + # DownRes 3 + conv_32 = self.double_conv_layer3(pool_32) # (32,32,32) + pool_16 = self.MaxPooling3(conv_32) # (32,16,16) + # DownRes 4 + conv_16 = self.double_conv_layer4(pool_16) # (64,16,16) + pool_8 = self.MaxPooling4(conv_16) # (64,8,8) + # DownRes 5, convolution only + conv_8 = self.double_conv_layer5(pool_8) # (128,8,8) + + # merge the other features, including Re, u_scaling and latent_z + if self.merge: + # merge_inputs # (8,8,1) + conv_8 = self.Concat([conv_8, merge_inputs]) # (129,8,8) + if self.latent_z: + # merge_inputs # (8,8,1) + # z_inputs # (8,8,50) + conv_8 = self.Concat([conv_8, merge_inputs, z_inputs]) # (8,8,179) + + # Up-sampling layers + # UpRes 6, attention gated concatenation + upsampling + double residual convolution + gating_16 = self.gating_signal1(conv_8) + up_16 = self.Conv2DTranspose1(conv_8) # (None, 128, 16, 16) + + # add channel attention block after spatial attention block + att_16 = self.attention_block1(conv_16, gating_16) # (None, 64, 16, 16) + up_16 = self.Concat([up_16, att_16]) # (None, 192, 16, 16) + se_conv_16 = self.se_conv_16_layer(up_16) # (None, 192, 16, 16) + up_conv_16 = self.up_conv_16_layer(se_conv_16) # (None, 64, 16, 16) + + # UpRes 7 + gating_32 = self.gating_signal2(up_conv_16) + up_32 = self.Conv2DTranspose2(up_conv_16) # (None, 64, 32, 32) + + att_32 = self.attention_block2(conv_32, gating_32) # (None, 32, 32, 32) + up_32 = self.Concat([up_32, att_32]) # (None, 96, 32, 32) + se_conv_32 = self.se_conv_32_layer(up_32) # (None, 96, 32, 32) + up_conv_32 = self.up_conv_32_layer(se_conv_32) # (None, 32, 32, 32) + + # UpRes 8 + gating_64 = self.gating_signal3(up_conv_32) + up_64 = self.Conv2DTranspose3(up_conv_32) # (None, 32, 64, 64) + + att_64 = self.attention_block3(conv_64, gating_64) # (None, 64, 64, 16) + up_64 = self.Concat([up_64, att_64]) # (None, 48, 64, 64) + se_conv_64 = self.se_conv_64_layer(up_64) # (None, 48, 64, 64) + up_conv_64 = self.up_conv_64_layer(se_conv_64) # (None, 16, 64, 64) + + # UpRes 9 + gating_128 = self.gating_signal4(up_conv_64) + up_128 = self.Conv2DTranspose4(up_conv_64) # (None, 16, 128, 128) + + att_128 = self.attention_block4(conv_128, gating_128) # (None, 8, 128, 128) + up_128 = self.Concat([up_128, att_128]) # (None, 24, 128, 128) + se_conv_128 = self.se_conv_128_layer(up_128) # (None, 24, 128, 128) + up_conv_128 = self.up_conv_128_layer(se_conv_128) # (None, 8, 128, 128) + + # 1*1 convolutional layers + # valid padding + # batch normalization + # sigmoid nonlinear activation + conv_final = self.conv_final_layer(up_conv_128) + conv_final = self.BatchNormal(conv_final) + conv_final = self.Tanh(conv_final) + return conv_final diff --git a/MindFlow/applications/research/cascade_net/src/dataset.py b/MindFlow/applications/research/cascade_net/src/dataset.py new file mode 100644 index 0000000000000000000000000000000000000000..0064d49751d23792a6748f14fc30659776613311 --- /dev/null +++ b/MindFlow/applications/research/cascade_net/src/dataset.py @@ -0,0 +1,61 @@ +# Copyright 2023 Huawei Technologies Co., Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +import scipy.io as sio +import numpy as np + + +class AccesstrainDataset: + def __init__(self, path): + self.u_r10_train = sio.loadmat(path + 'ur_10.mat')['ur_10_train'].astype(np.float32) # (10792, 128, 128, 2) + self.u_r5_train = sio.loadmat(path + 'ur_5.mat')['ur_5_train'].astype(np.float32) # (10792, 128, 128, 2) + self.u_r3_train = sio.loadmat(path + 'ur_3.mat')['ur_3_train'].astype(np.float32) # (10792, 128, 128, 2) + self.u_r1_train = sio.loadmat(path + 'ur_1.mat')['ur_1_train'].astype(np.float32) # (10792, 128, 128, 2) + self.cp_fluc_train = sio.loadmat(path + 'Input_flucReal_Cp_pd.mat')['cp_flucReal_train_dp'].astype( + np.float32) # (10792, 128, 128, 3) + self.Re_c_train = sio.loadmat(path + 'conditions_Re_pd.mat')['Re_c_train'].astype(np.float32) # (10792, 1) + self.scaling_input_train = sio.loadmat(path + 'Input_scaling_u_pd.mat')['fine_input_train'].astype( + np.float32) # (10792,20) + + def __getitem__(self, index): + return (self.u_r10_train[index], self.u_r5_train[index], self.u_r3_train[index], self.u_r1_train[index], + self.cp_fluc_train[index], self.Re_c_train[index], self.scaling_input_train[index],) + + def __len__(self): + return len(self.u_r10_train) + + +def validation_test_dataset(path): + u_r10_validation = sio.loadmat(path + 'ur_10.mat')['ur_10_validation'].astype(np.float32) # (1704, 128, 128, 2) + u_r5_validation = sio.loadmat(path + 'ur_5.mat')['ur_5_validation'].astype(np.float32) # (1704, 128, 128, 2) + u_r3_validation = sio.loadmat(path + 'ur_3.mat')['ur_3_validation'].astype(np.float32) # (1704, 128, 128, 2) + u_r1_validation = sio.loadmat(path + 'ur_1.mat')['ur_1_validation'].astype(np.float32) # (1704, 128, 128, 2) + cp_fluc_validation = sio.loadmat(path + 'Input_flucReal_Cp_pd.mat')['cp_flucReal_validation_dp'].astype( + np.float32) # (1704, 128, 128, 3) + Re_c_validation = sio.loadmat(path + 'conditions_Re_pd.mat')['Re_c_validation'].astype(np.float32) # (1704,1) + scaling_input_validation = sio.loadmat(path + 'Input_scaling_u_pd.mat')['fine_input_validation'].astype( + np.float32) # (1704, 20) + + u_r10_test = sio.loadmat(path + 'ur_10.mat')['ur_10_test'].astype(np.float32) # (3408, 128, 128, 2) + u_r5_test = sio.loadmat(path + 'ur_5.mat')['ur_5_test'].astype(np.float32) # (3408, 128, 128, 2) + u_r3_test = sio.loadmat(path + 'ur_3.mat')['ur_3_test'].astype(np.float32) # (3408, 128, 128, 2) + u_r1_test = sio.loadmat(path + 'ur_1.mat')['ur_1_test'].astype(np.float32) # (3408, 128, 128, 2) + cp_fluc_test = sio.loadmat(path + 'Input_flucReal_Cp_pd.mat')['cp_flucReal_test_dp'].astype( + np.float32) # (3408, 128, 128, 3) + Re_c_test = sio.loadmat(path + 'conditions_Re_pd.mat')['Re_c_test'].astype(np.float32) # (3408, 1) + scaling_input_test = sio.loadmat(path + 'Input_scaling_u_pd.mat')['fine_input_test'].astype(np.float32) # (3408,20) + + return (u_r10_validation, u_r5_validation, u_r3_validation, u_r1_validation, + cp_fluc_validation, Re_c_validation, scaling_input_validation, + u_r10_test, u_r5_test, u_r3_test, u_r1_test, cp_fluc_test, Re_c_test, scaling_input_test) \ No newline at end of file diff --git a/MindFlow/applications/research/cascade_net/train.py b/MindFlow/applications/research/cascade_net/train.py new file mode 100644 index 0000000000000000000000000000000000000000..b4976fa992fb5c64e92096249e96440b8f5f6e8e --- /dev/null +++ b/MindFlow/applications/research/cascade_net/train.py @@ -0,0 +1,358 @@ +# Copyright 2023 Huawei Technologies Co., Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +import gc +import argparse + +from mindspore.dataset import GeneratorDataset +import mindspore as ms +from mindspore import ops, Tensor +import mindspore.nn as nn +import numpy as np +import matplotlib.pyplot as plt +import time +from tqdm import tqdm +import matplotlib +from mindflow.utils import load_yaml_config +from mindspore import context +from src.Compound_Model import init_sub_model, DefineCompoundCritic, DefineCompoundGan +from src.Loss import WassersteinLoss, GradLoss +from src.dataset import AccesstrainDataset, validation_test_dataset +matplotlib.use('agg') # to prevent "Fail to create pixmap with Tk_GetPixmap in TkImgPhotoInstanceSetSize" +ms.dataset.config.set_prefetch_size(1) +context.set_context(device_target="GPU", mode=context.PYNATIVE_MODE) + +# ---------------------------------source locations------------------------------------ +path = 'dataset/' + +# 数据读取 +loader = AccesstrainDataset(path) +train_dataset = GeneratorDataset(source=loader, column_names=[ + 'u_r10_train', 'u_r5_train', 'u_r3_train', 'u_r1_train', 'cp_fluc_train', 'Re_c_train', 'scaling_input_train']) +train_dataset = train_dataset.shuffle(buffer_size=25) + +(u_r10_validation, u_r5_validation, u_r3_validation, u_r1_validation, cp_fluc_validation, Re_c_validation, + scaling_input_validation, u_r10_test, u_r5_test, u_r3_test, u_r1_test, cp_fluc_test, Re_c_test, scaling_input_test) \ + = validation_test_dataset(path) + + +def train(train_dataset, config): + # 参数 + sample_interval = config["sample_interval"] + batch_size = config["batch_size"] + train_dataset = train_dataset.batch(batch_size, drop_remainder=True) + n_channel_p = config["n_channel_p"] + n_channel_u = config["n_channel_u"] + lambda_GP = config["lambda_GP"] + lambda_L2_u = config["lambda_L2_u"] + lambda_L2_gradu = config["lambda_L2_gradu"] + critic_model_lr = config["critic_model_lr"] + gan_model_lr = config["gan_model_lr"] + n_critic = config["n_critic"] + sample_num_train = loader.__len__() + sample_num_val = u_r1_validation.shape[0] + sample_num_test = u_r1_test.shape[0] + plot_n = config["plot_n"] + n_imgs = config["n_imgs"] + merge_n_imgs = config["merge_n_imgs"] + latent_z_n_channel = config["latent_z_n_channel"] + epochs = config["epochs"] + losslog = [] + + # 生成器及判别器 + (merge_model, g_model_I, d_model_I, g_model_II, d_model_II, + g_model_III, d_model_III, g_model_IV, d_model_IV) = init_sub_model(n_channel_p, n_channel_u) + critic_model = DefineCompoundCritic(n_channel_p, n_channel_u, batch_size, + d_model_I, d_model_II, d_model_III, d_model_IV) + critic_model.update_parameters_name('critic') + gan_model = DefineCompoundGan(n_channel_p, n_channel_u, merge_model, + g_model_I, g_model_II, g_model_III, g_model_IV) + gan_model.update_parameters_name('generator') + + # 损失函数 + wasserstein_loss = WassersteinLoss() + mae_loss = nn.MAELoss(reduction='none') + grad_loss = GradLoss() + + # 前向传播 + def d_forward_fn(input): + g_input = [input[4], input[5], input[6], input[7]] + g_pred = gan_model(g_input) + d_input = [input[0], input[1], input[2], input[3], g_pred[4], g_pred[0], g_pred[1], g_pred[2], g_pred[3]] + + d_pred = critic_model(d_input) + loss_wass = wasserstein_loss(ops.stack(d_pred[:8]), ops.Concat(axis=0)( + [-ops.ones_like(ops.stack(d_pred[:4])), ops.ones_like(ops.stack(d_pred[4:8]))])) + loss_gradp = ops.stack([d_pred[12](), d_pred[13](), d_pred[14](), d_pred[15]()]) + loss = ops.sum(1 * loss_wass) + ops.sum(lambda_GP * loss_gradp) + loss_list = ms.numpy.concatenate((ms.numpy.expand_dims(loss, 0), loss_wass, loss_gradp), axis=0).asnumpy() + return loss, loss_list + + def g_forward_fn(true, input): + g_input = [input[4], input[5], input[6], input[7]] + g_pred = gan_model(g_input) + d_input = [input[0], input[1], input[2], input[3], g_pred[4], g_pred[0], g_pred[1], g_pred[2], g_pred[3]] + d_pred = critic_model(d_input) + + loss_wass = wasserstein_loss(ops.stack(d_pred[4:8]), -ops.ones_like(ops.stack(d_pred[4:8]))) + loss_mae = mae_loss(ops.stack(g_pred[:4]), ops.stack(true[:4])) + loss_mae = ops.mean(loss_mae, axis=(1, 2, 3, 4)) + loss_grad = grad_loss(ops.stack(g_pred[5:9]), ops.stack(true[:4])) + loss = (ops.sum(1 * loss_wass) + ops.sum(lambda_L2_u * loss_mae) + ops.sum(lambda_L2_gradu * loss_grad)) + loss_list = ms.numpy.concatenate((ms.numpy.expand_dims(loss, 0), loss_wass, loss_mae, loss_grad), axis=0).asnumpy() + return loss, loss_list + + # 优化器 + d_optimizer = nn.RMSProp(critic_model.trainable_params(), learning_rate=critic_model_lr, epsilon=1e-07) + d_grad_fn = ms.value_and_grad(d_forward_fn, None, d_optimizer.parameters, has_aux=True) + g_optimizer = nn.RMSProp(gan_model.trainable_params(), learning_rate=gan_model_lr, epsilon=1e-07) + g_grad_fn = ms.value_and_grad(g_forward_fn, None, g_optimizer.parameters, has_aux=True) + d_optimizer.update_parameters_name('optim_d') + g_optimizer.update_parameters_name('optim_g') + + def train_step(g_real_data, input): + # 计算判别器损失和梯度 + for _ in range(n_critic): + (_, d_loss_list), d_grads = d_grad_fn(input) + d_optimizer(d_grads) + ave_g_loss_train_ncritic.append(d_loss_list) + ave_d_loss_train.append(np.mean(ave_g_loss_train_ncritic, axis=0)) + (_, g_loss_list), g_grads = g_grad_fn(g_real_data, input) + g_optimizer(g_grads) + ave_g_loss_train.append(g_loss_list) + + def define_evaluation(u_r10, u_r5, u_r3, u_r1, cp, Re, scaling_input, latent_z, idx): + _, batch_d_loss = d_forward_fn( + [Tensor(u_r10[idx, :, :, :]), Tensor(u_r5[idx, :, :, :]), Tensor(u_r3[idx, :, :, :]), Tensor(u_r1[idx, :, :, :]), + Tensor(cp[idx, :, :, :]), Tensor(Re[idx, :]), Tensor(scaling_input[idx, :]), latent_z]) + _, batch_g_loss = g_forward_fn( + [Tensor(u_r10[idx, :, :, :]), Tensor(u_r5[idx, :, :, :]), Tensor(u_r3[idx, :, :, :]), Tensor(u_r1[idx, :, :, :])], + [Tensor(u_r10[idx, :, :, :]), Tensor(u_r5[idx, :, :, :]), Tensor(u_r3[idx, :, :, :]), Tensor(u_r1[idx, :, :, :]), + Tensor(cp[idx, :, :, :]), Tensor(Re[idx, :]), Tensor(scaling_input[idx, :]), latent_z]) + return batch_d_loss, batch_g_loss + + def sample_images(epoch): # plot_n = 5 + idx_train = np.random.randint(0, sample_num_train, plot_n) + idx_test = np.random.randint(0, sample_num_test, plot_n) + latent_z_input_train = Tensor(np.random.normal(0, 1, ( + plot_n, latent_z_n_channel, merge_n_imgs, merge_n_imgs)), dtype=ms.float32) + latent_z_input_test = Tensor(np.random.normal(0, 1, ( + plot_n, latent_z_n_channel, merge_n_imgs, merge_n_imgs)), dtype=ms.float32) + gen_u_I, gen_u_II, gen_u_III, gen_u_IV, _, _, _, _, _, _ = gan_model([ + Tensor(loader[idx_train][4]), + Tensor(loader[idx_train][5]), + Tensor(loader[idx_train][6]), + latent_z_input_train]) + gen_u_I_, gen_u_II_, gen_u_III_, gen_u_IV_, _, _, _, _, _, _ = gan_model([ + Tensor(cp_fluc_test[idx_test, :, :, :]), + Tensor(Re_c_test[idx_test, :]), + Tensor(scaling_input_test[idx_test, :]), + latent_z_input_test]) + del _ + plot_images(epoch, + velocity_part=0, + gen_u_I=gen_u_I, + gen_u_II=gen_u_II, + gen_u_III=gen_u_III, + gen_u_IV=gen_u_IV, + sel_index=idx_train, + is_train=True) + plot_images(epoch, + velocity_part=1, + gen_u_I=gen_u_I, + gen_u_II=gen_u_II, + gen_u_III=gen_u_III, + gen_u_IV=gen_u_IV, + sel_index=idx_train, + is_train=True) + plot_images(epoch, + velocity_part=0, + gen_u_I=gen_u_I_, + gen_u_II=gen_u_II_, + gen_u_III=gen_u_III_, + gen_u_IV=gen_u_IV_, + sel_index=idx_test, + is_test=True) + plot_images(epoch, + velocity_part=1, + gen_u_I=gen_u_I_, + gen_u_II=gen_u_II_, + gen_u_III=gen_u_III_, + gen_u_IV=gen_u_IV_, + sel_index=idx_test, is_test=True) + + def plot_images(epoch, velocity_part, gen_u_I, gen_u_II, gen_u_III, gen_u_IV, sel_index, is_train=False, + is_test=False): + # including ur,vr,gradur, and grad_vr + # velocity_part = 0 for u and 1 for v + velocity_name = list(['u', 'v']) + fig, ax = plt.subplots(5, 8, figsize=(20, 16)) + if is_train: + index = 0 + for i in range(0, 5, 1): + ax[i, 0].contourf(xx, yy, + gen_u_I[index, velocity_part, :, :].reshape(n_imgs, n_imgs).asnumpy(), + cmap='coolwarm') + ax[i, 0].axis('off') + ax[i, 1].contourf(xx, yy, loader[sel_index[index]][0][velocity_part, :, :], cmap='coolwarm') + ax[i, 1].axis('off') + ax[i, 2].contourf(xx, yy, + gen_u_II[index, velocity_part, :, :].reshape(n_imgs, n_imgs).asnumpy(), + cmap='coolwarm') + ax[i, 2].axis('off') + ax[i, 3].contourf(xx, yy, loader[sel_index[index]][1][velocity_part, :, :], cmap='coolwarm') + ax[i, 3].axis('off') + ax[i, 4].contourf(xx, yy, + gen_u_III[index, velocity_part, :, :].reshape(n_imgs, n_imgs).asnumpy(), + cmap='coolwarm') + ax[i, 4].axis('off') + ax[i, 5].contourf(xx, yy, loader[sel_index[index]][2][velocity_part, :, :], cmap='coolwarm') + ax[i, 5].axis('off') + ax[i, 6].contourf(xx, yy, + gen_u_IV[index, velocity_part, :, :].reshape(n_imgs, n_imgs).asnumpy(), + cmap='coolwarm') + ax[i, 6].axis('off') + ax[i, 7].contourf(xx, yy, loader[sel_index[index]][3][velocity_part, :, :], cmap='coolwarm') + ax[i, 7].axis('off') + + index = index + 1 + plt.savefig('training_results/images/imgs_train_' + velocity_name[velocity_part] + '%d.png' % ( + epoch)) + plt.close() + if is_test: + index = 0 + for i in range(0, 5, 1): + ax[i, 0].contourf(xx, yy, + gen_u_I[index, velocity_part, :, :].reshape(n_imgs, n_imgs).asnumpy(), + cmap='coolwarm') + ax[i, 0].axis('off') + ax[i, 1].contourf(xx, yy, + u_r10_test[sel_index[index], velocity_part, :, :].reshape(n_imgs, n_imgs), + cmap='coolwarm') + ax[i, 1].axis('off') + ax[i, 2].contourf(xx, yy, + gen_u_II[index, velocity_part, :, :].reshape(n_imgs, n_imgs).asnumpy(), + cmap='coolwarm') + ax[i, 2].axis('off') + ax[i, 3].contourf(xx, yy, + u_r5_test[sel_index[index], velocity_part, :, :].reshape(n_imgs, n_imgs), + cmap='coolwarm') + ax[i, 3].axis('off') + ax[i, 4].contourf(xx, yy, + gen_u_III[index, velocity_part, :, :].reshape(n_imgs, n_imgs).asnumpy(), + cmap='coolwarm') + ax[i, 4].axis('off') + ax[i, 5].contourf(xx, yy, + u_r3_test[sel_index[index], velocity_part, :, :].reshape(n_imgs, n_imgs), + cmap='coolwarm') + ax[i, 5].axis('off') + ax[i, 6].contourf(xx, yy, + gen_u_IV[index, velocity_part, :, :].reshape(n_imgs, n_imgs).asnumpy(), + cmap='coolwarm') + ax[i, 6].axis('off') + ax[i, 7].contourf(xx, yy, + u_r1_test[sel_index[index], velocity_part, :, :].reshape(n_imgs, n_imgs), + cmap='coolwarm') + ax[i, 7].axis('off') + index = index + 1 + plt.savefig('training_results/images/imgs_test_' + velocity_name[velocity_part] + '%d.png' % ( + epoch)) + plt.close() + plt.cla() + plt.close('all') + + critic_model.set_train() + gan_model.set_train() + + y1 = np.linspace(-1, 1, 128) + x1 = np.linspace(-1, 1, 128) + xx, yy = np.meshgrid(x1, y1) + + for epoch in range(epochs): + start = time.perf_counter() + ave_d_loss_train = list() + ave_g_loss_train = list() + for _, (u_r10_train, u_r5_train, u_r3_train, u_r1_train, cp_fluc_train, Re_c_train, scaling_input_train) in ( + enumerate(tqdm(train_dataset))): + ave_g_loss_train_ncritic = list() + latent_z_input_train = ops.normal( + (batch_size, latent_z_n_channel, merge_n_imgs, merge_n_imgs), + 0, 1, ) + input = [u_r10_train, u_r5_train, u_r3_train, u_r1_train, + cp_fluc_train, Re_c_train, scaling_input_train, latent_z_input_train] + g_true = [u_r10_train, u_r5_train, u_r3_train, u_r1_train] + train_step(g_true, input) + ave_d_loss_train = np.mean(ave_d_loss_train, axis=0) + ave_g_loss_train = np.mean(ave_g_loss_train, axis=0) + + # predict the d_loss and g_loss in validation + idx_ = np.random.randint(0, sample_num_val, batch_size) + latent_z_input_validation = Tensor(np.random.normal(0, 1, ( + batch_size, latent_z_n_channel, merge_n_imgs, merge_n_imgs)), dtype=ms.float32) + batch_d_loss_val, batch_g_loss_val = define_evaluation( + u_r10_validation, u_r5_validation, u_r3_validation, u_r1_validation, + cp_fluc_validation, Re_c_validation, scaling_input_validation, latent_z_input_validation, idx_) + + # predict the d_loss and g_loss in testing + idx__ = np.random.randint(0, sample_num_test, batch_size) + latent_z_input_test = Tensor(np.random.normal(0, 1, ( + batch_size, latent_z_n_channel, merge_n_imgs, merge_n_imgs)), dtype=ms.float32) + batch_d_loss_test, batch_g_loss_test = define_evaluation( + u_r10_test, u_r5_test, u_r3_test, u_r1_test, + cp_fluc_test, Re_c_test, scaling_input_test, latent_z_input_test, idx__) + + end = time.perf_counter() + time_consumption = (end - start) // 60 # covert s. to min. + # Plot the progress and index[0] is the weighted loss + print( + "%d:[%d min] [D:%.2f, %.2f, %.2f, %.2f, %.2f] [G:%.2f, %.2f, %.2f, %.2f, %.2f] " + "[L2_train:%.5f, %.5f, %.5f, %.5f] [L2_val:%.5f, %.5f, %.5f, %.5f] [L2_test:%.5f, %.5f, %.5f, %.5f]" % + (epoch, time_consumption, + ave_d_loss_train[0], + ave_d_loss_train[1] + ave_d_loss_train[5] + ave_d_loss_train[9], + ave_d_loss_train[2] + ave_d_loss_train[6] + ave_d_loss_train[10], + ave_d_loss_train[3] + ave_d_loss_train[7] + ave_d_loss_train[11], + ave_d_loss_train[4] + ave_d_loss_train[8] + ave_d_loss_train[12], + ave_g_loss_train[0], + ave_g_loss_train[1], ave_g_loss_train[2], ave_g_loss_train[3], ave_g_loss_train[4], + ave_g_loss_train[5], ave_g_loss_train[6], ave_g_loss_train[7], ave_g_loss_train[8], + batch_g_loss_val[5], batch_g_loss_val[6], batch_g_loss_val[7], batch_g_loss_val[8], + batch_g_loss_test[5], batch_g_loss_test[6], batch_g_loss_test[7], batch_g_loss_test[8])) + + losslog.append(np.concatenate( # 60 + ([ave_d_loss_train[i] for i in range(13)], # 13 + [batch_d_loss_val[i] for i in range(13)], # 13 + [batch_d_loss_test[i] for i in range(13)], # 13 + [ave_g_loss_train[i] for i in range(13)], # 13 + [batch_g_loss_val[i] for i in range(13)], # 13 + [batch_g_loss_test[i] for i in range(13)]))) # 13 + + # If at save interval => save generated image samples + if epoch % sample_interval == 0: + sample_images(epoch) + ms.save_checkpoint(gan_model, "training_results/model/gan_%d" % epoch + ".ckpt") + ms.save_checkpoint(critic_model, "training_results/model/critic_%d" % epoch + ".ckpt") + np.savetxt('training_results/loss/loss.txt', losslog, fmt='%.4f') + del ave_d_loss_train, ave_g_loss_train, batch_d_loss_val, batch_g_loss_val, batch_d_loss_test, batch_g_loss_test + gc.collect() + + + +if __name__ == '__main__': + parser = argparse.ArgumentParser() + parser.add_argument("--config_path", type=str, default="./config/Cascade-Net.yaml") + args = parser.parse_args() + config = load_yaml_config(args.config_path) + + train(train_dataset, config) diff --git a/MindFlow/applications/research/cascade_net/training_results/images/.keep b/MindFlow/applications/research/cascade_net/training_results/images/.keep new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/MindFlow/applications/research/cascade_net/training_results/loss/.keep b/MindFlow/applications/research/cascade_net/training_results/loss/.keep new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/MindFlow/applications/research/cascade_net/training_results/model/.keep b/MindFlow/applications/research/cascade_net/training_results/model/.keep new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391