NowcastNet: Physics-based generative model for extreme precipitation nowcasting
Overview
NowcastNet is proposed by Long Mingsheng. It is a nonlinear nowcasting model for extreme precipitation with lead times of up to 3h. It includes a stochastic generative network and a deterministic evolution network. The evolution network yields physically plausible predictions for advective features at a scale of 20km. The generative network generates convective details present in the radar observations. More information can be found in paper. The architecture of the Nowcastnet is shown below.
NowcastNet
Evolution network：a motion decoder for learning motion fields \(v_{1:T}\) and an intensity decoder for learning intensity residuals \(s_{1:T}\). Then, evolution operator uses \(v_{1:T}\) and \(s_{1:T}\) to predict \(x_{1:T}^{''}\). The formalization is shown below.
\[x_{1:T}^{''} = Evolution(x_{-T:0})\]
Nowcast encoder & decoder：It uses Semantic Image Synthesis with Spatially-Adaptive Normalization to model the relation between physics-conditioning \(x_{1:T}^{''}\) and outputs.
Technology path
MindSpore Earth solves the problem as follows:
Data Construction.
Model Construction.
Loss function.
Model Training.
Model Evaluation and Visualization.
Download the training and test dataset: Nowcastnet/dataset
[2]:
import random
import mindspore as ms
import numpy as np
from mindspore import context, nn, amp, set_seed, load_checkpoint, load_param_into_net
[17]:
from src import get_logger
from src import EvolutionTrainer, GenerationTrainer, GenerateLoss, DiscriminatorLoss, EvolutionLoss
from src import EvolutionPredictor, GenerationPredictor
from src import RadarData, NowcastDataset
from src.evolution import EvolutionNet
from src.generator import GenerationNet
from src.discriminator import TemporalDiscriminator
from src.visual import plt_img
from mindearth.utils.tools import load_yaml_config
[4]:
np.random.seed(0)
set_seed(0)
random.seed(0)
[5]:
config = load_yaml_config("./configs/Nowcastnet.yaml")
context.set_context(mode=context.GRAPH_MODE, device_target="Ascend", device_id=1)
Data Construction
Download the statistic, training and validation dataset from dataset to ./dataset. Modify the parameter of root_dir in the Nowcastnet.yaml, which set the directory for dataset.
The ./dataset is hosted with the following directory structure:
├── train
├── valid
├── test
Evolution
[6]:
logger = get_logger(config)
config["model"]["module_name"] = 'evolution'
config["data"]["batch_size"] = 4
config["summary"]["eval_interval"] = 1
config["summary"]["visual"] = False
train_params = config.get("train")
summary_params = config.get("summary")
evo_model = EvolutionNet(config)
evo_model.set_train()
2024-01-30 20:21:45,539 - utils.py[line:55] - INFO: {'name': 'NowcastNet', 'ngf': 32, 'pool_ensemble_num': 4, 'module_name': 'generation'}
2024-01-30 20:21:45,539 - utils.py[line:55] - INFO: {'name': 'us', 'root_dir': './datasets', 't_in': 9, 't_out': 20, 'h_size': 512, 'w_size': 512, 'time_interval': 10, 'num_workers': 1, 'data_sink': False, 'batch_size': 1, 'noise_scale': 32}
2024-01-30 20:21:45,540 - utils.py[line:55] - INFO: {'name': 'adam', 'beta1': 0.01, 'beta2': 0.9, 'g_lr': 1.5e-05, 'd_lr': '6e-5', 'epochs': 10}
2024-01-30 20:21:45,541 - utils.py[line:55] - INFO: {'name': 'adam', 'lr': 0.001, 'weight_decay': 0.1, 'gamma': 0.5, 'epochs': 5}
2024-01-30 20:21:45,541 - utils.py[line:55] - INFO: {'summary_dir': './summary/', 'eval_interval': 2, 'save_checkpoint_epochs': 2, 'keep_checkpoint_max': 4, 'key_info_timestep': [10, 60, 120], 'generate_ckpt_path': '', 'evolution_ckpt_path': '', 'visual': True, 'csin_threshold': 16}
2024-01-30 20:21:45,542 - utils.py[line:55] - INFO: {'distribute': False, 'mixed_precision': True, 'amp_level': 'O2', 'load_ckpt': False}
[6]:
EvolutionNet<
(evo_net): EvolutionNetwork<
(inc): DoubleConv<
(single_conv): SequentialCell<
(0): BatchNorm2d<num_features=9, eps=1e-05, momentum=0.9, gamma=Parameter (name=evo_net.inc.single_conv.0.gamma, shape=(9,), dtype=Float32, requires_grad=True), beta=Parameter (name=evo_net.inc.single_conv.0.beta, shape=(9,), dtype=Float32, requires_grad=True), moving_mean=Parameter (name=evo_net.inc.single_conv.0.moving_mean, shape=(9,), dtype=Float32, requires_grad=False), moving_variance=Parameter (name=evo_net.inc.single_conv.0.moving_variance, shape=(9,), dtype=Float32, requires_grad=False)>
(1): SpectralNormal<
(parametrizations): Conv2d<input_channels=9, output_channels=32, kernel_size=(3, 3), stride=(1, 1), pad_mode=pad, padding=1, dilation=(1, 1), group=1, has_bias=True, weight_init=<mindspore.common.initializer.HeUniform object at 0xfffe7c6c2df0>, bias_init=<mindspore.common.initializer.Uniform object at 0xfffe777e12b0>, format=NCHW>
>
>
(double_conv): SequentialCell<
(0): BatchNorm2d<num_features=9, eps=1e-05, momentum=0.9, gamma=Parameter (name=evo_net.inc.double_conv.0.gamma, shape=(9,), dtype=Float32, requires_grad=True), beta=Parameter (name=evo_net.inc.double_conv.0.beta, shape=(9,), dtype=Float32, requires_grad=True), moving_mean=Parameter (name=evo_net.inc.double_conv.0.moving_mean, shape=(9,), dtype=Float32, requires_grad=False), moving_variance=Parameter (name=evo_net.inc.double_conv.0.moving_variance, shape=(9,), dtype=Float32, requires_grad=False)>
(1): ReLU<>
(2): SpectralNormal<
(parametrizations): Conv2d<input_channels=9, output_channels=32, kernel_size=(3, 3), stride=(1, 1), pad_mode=pad, padding=1, dilation=(1, 1), group=1, has_bias=True, weight_init=<mindspore.common.initializer.HeUniform object at 0xfffe777f1d90>, bias_init=<mindspore.common.initializer.Uniform object at 0xfffe777f1f10>, format=NCHW>
>
(3): BatchNorm2d<num_features=32, eps=1e-05, momentum=0.9, gamma=Parameter (name=evo_net.inc.double_conv.3.gamma, shape=(32,), dtype=Float32, requires_grad=True), beta=Parameter (name=evo_net.inc.double_conv.3.beta, shape=(32,), dtype=Float32, requires_grad=True), moving_mean=Parameter (name=evo_net.inc.double_conv.3.moving_mean, shape=(32,), dtype=Float32, requires_grad=False), moving_variance=Parameter (name=evo_net.inc.double_conv.3.moving_variance, shape=(32,), dtype=Float32, requires_grad=False)>
(4): ReLU<>
(5): SpectralNormal<
(parametrizations): Conv2d<input_channels=32, output_channels=32, kernel_size=(3, 3), stride=(1, 1), pad_mode=pad, padding=1, dilation=(1, 1), group=1, has_bias=True, weight_init=<mindspore.common.initializer.HeUniform object at 0xfffe74306640>, bias_init=<mindspore.common.initializer.Uniform object at 0xfffe743067c0>, format=NCHW>
>
>
>
(down1): Down<
(maxpool_conv): SequentialCell<
(0): MaxPool2d<kernel_size=2, stride=2, pad_mode=VALID>
(1): DoubleConv<
(single_conv): SequentialCell<
(0): BatchNorm2d<num_features=32, eps=1e-05, momentum=0.9, gamma=Parameter (name=evo_net.down1.maxpool_conv.1.single_conv.0.gamma, shape=(32,), dtype=Float32, requires_grad=True), beta=Parameter (name=evo_net.down1.maxpool_conv.1.single_conv.0.beta, shape=(32,), dtype=Float32, requires_grad=True), moving_mean=Parameter (name=evo_net.down1.maxpool_conv.1.single_conv.0.moving_mean, shape=(32,), dtype=Float32, requires_grad=False), moving_variance=Parameter (name=evo_net.down1.maxpool_conv.1.single_conv.0.moving_variance, shape=(32,), dtype=Float32, requires_grad=False)>
(1): SpectralNormal<
(parametrizations): Conv2d<input_channels=32, output_channels=64, kernel_size=(3, 3), stride=(1, 1), pad_mode=pad, padding=1, dilation=(1, 1), group=1, has_bias=True, weight_init=<mindspore.common.initializer.HeUniform object at 0xfffe74306a30>, bias_init=<mindspore.common.initializer.Uniform object at 0xfffe74306a60>, format=NCHW>
>
>
(double_conv): SequentialCell<
(0): BatchNorm2d<num_features=32, eps=1e-05, momentum=0.9, gamma=Parameter (name=evo_net.down1.maxpool_conv.1.double_conv.0.gamma, shape=(32,), dtype=Float32, requires_grad=True), beta=Parameter (name=evo_net.down1.maxpool_conv.1.double_conv.0.beta, shape=(32,), dtype=Float32, requires_grad=True), moving_mean=Parameter (name=evo_net.down1.maxpool_conv.1.double_conv.0.moving_mean, shape=(32,), dtype=Float32, requires_grad=False), moving_variance=Parameter (name=evo_net.down1.maxpool_conv.1.double_conv.0.moving_variance, shape=(32,), dtype=Float32, requires_grad=False)>
(1): ReLU<>
(2): SpectralNormal<
(parametrizations): Conv2d<input_channels=32, output_channels=64, kernel_size=(3, 3), stride=(1, 1), pad_mode=pad, padding=1, dilation=(1, 1), group=1, has_bias=True, weight_init=<mindspore.common.initializer.HeUniform object at 0xfffe74306c10>, bias_init=<mindspore.common.initializer.Uniform object at 0xfffe7cec2e20>, format=NCHW>
>
(3): BatchNorm2d<num_features=64, eps=1e-05, momentum=0.9, gamma=Parameter (name=evo_net.down1.maxpool_conv.1.double_conv.3.gamma, shape=(64,), dtype=Float32, requires_grad=True), beta=Parameter (name=evo_net.down1.maxpool_conv.1.double_conv.3.beta, shape=(64,), dtype=Float32, requires_grad=True), moving_mean=Parameter (name=evo_net.down1.maxpool_conv.1.double_conv.3.moving_mean, shape=(64,), dtype=Float32, requires_grad=False), moving_variance=Parameter (name=evo_net.down1.maxpool_conv.1.double_conv.3.moving_variance, shape=(64,), dtype=Float32, requires_grad=False)>
(4): ReLU<>
(5): SpectralNormal<
(parametrizations): Conv2d<input_channels=64, output_channels=64, kernel_size=(3, 3), stride=(1, 1), pad_mode=pad, padding=1, dilation=(1, 1), group=1, has_bias=True, weight_init=<mindspore.common.initializer.HeUniform object at 0xfffe74306e80>, bias_init=<mindspore.common.initializer.Uniform object at 0xfffe74306f40>, format=NCHW>
>
>
>
>
>
(down2): Down<
(maxpool_conv): SequentialCell<
(0): MaxPool2d<kernel_size=2, stride=2, pad_mode=VALID>
(1): DoubleConv<
(single_conv): SequentialCell<
(0): BatchNorm2d<num_features=64, eps=1e-05, momentum=0.9, gamma=Parameter (name=evo_net.down2.maxpool_conv.1.single_conv.0.gamma, shape=(64,), dtype=Float32, requires_grad=True), beta=Parameter (name=evo_net.down2.maxpool_conv.1.single_conv.0.beta, shape=(64,), dtype=Float32, requires_grad=True), moving_mean=Parameter (name=evo_net.down2.maxpool_conv.1.single_conv.0.moving_mean, shape=(64,), dtype=Float32, requires_grad=False), moving_variance=Parameter (name=evo_net.down2.maxpool_conv.1.single_conv.0.moving_variance, shape=(64,), dtype=Float32, requires_grad=False)>
(1): SpectralNormal<
(parametrizations): Conv2d<input_channels=64, output_channels=128, kernel_size=(3, 3), stride=(1, 1), pad_mode=pad, padding=1, dilation=(1, 1), group=1, has_bias=True, weight_init=<mindspore.common.initializer.HeUniform object at 0xfffe742d1070>, bias_init=<mindspore.common.initializer.Uniform object at 0xfffe742d11c0>, format=NCHW>
>
>
(double_conv): SequentialCell<
(0): BatchNorm2d<num_features=64, eps=1e-05, momentum=0.9, gamma=Parameter (name=evo_net.down2.maxpool_conv.1.double_conv.0.gamma, shape=(64,), dtype=Float32, requires_grad=True), beta=Parameter (name=evo_net.down2.maxpool_conv.1.double_conv.0.beta, shape=(64,), dtype=Float32, requires_grad=True), moving_mean=Parameter (name=evo_net.down2.maxpool_conv.1.double_conv.0.moving_mean, shape=(64,), dtype=Float32, requires_grad=False), moving_variance=Parameter (name=evo_net.down2.maxpool_conv.1.double_conv.0.moving_variance, shape=(64,), dtype=Float32, requires_grad=False)>
(1): ReLU<>
(2): SpectralNormal<
(parametrizations): Conv2d<input_channels=64, output_channels=128, kernel_size=(3, 3), stride=(1, 1), pad_mode=pad, padding=1, dilation=(1, 1), group=1, has_bias=True, weight_init=<mindspore.common.initializer.HeUniform object at 0xfffe742d1490>, bias_init=<mindspore.common.initializer.Uniform object at 0xfffe742d1550>, format=NCHW>
>
(3): BatchNorm2d<num_features=128, eps=1e-05, momentum=0.9, gamma=Parameter (name=evo_net.down2.maxpool_conv.1.double_conv.3.gamma, shape=(128,), dtype=Float32, requires_grad=True), beta=Parameter (name=evo_net.down2.maxpool_conv.1.double_conv.3.beta, shape=(128,), dtype=Float32, requires_grad=True), moving_mean=Parameter (name=evo_net.down2.maxpool_conv.1.double_conv.3.moving_mean, shape=(128,), dtype=Float32, requires_grad=False), moving_variance=Parameter (name=evo_net.down2.maxpool_conv.1.double_conv.3.moving_variance, shape=(128,), dtype=Float32, requires_grad=False)>
(4): ReLU<>
(5): SpectralNormal<
(parametrizations): Conv2d<input_channels=128, output_channels=128, kernel_size=(3, 3), stride=(1, 1), pad_mode=pad, padding=1, dilation=(1, 1), group=1, has_bias=True, weight_init=<mindspore.common.initializer.HeUniform object at 0xfffe742d1310>, bias_init=<mindspore.common.initializer.Uniform object at 0xfffe742d1610>, format=NCHW>
>
>
>
>
>
(down3): Down<
(maxpool_conv): SequentialCell<
(0): MaxPool2d<kernel_size=2, stride=2, pad_mode=VALID>
(1): DoubleConv<
(single_conv): SequentialCell<
(0): BatchNorm2d<num_features=128, eps=1e-05, momentum=0.9, gamma=Parameter (name=evo_net.down3.maxpool_conv.1.single_conv.0.gamma, shape=(128,), dtype=Float32, requires_grad=True), beta=Parameter (name=evo_net.down3.maxpool_conv.1.single_conv.0.beta, shape=(128,), dtype=Float32, requires_grad=True), moving_mean=Parameter (name=evo_net.down3.maxpool_conv.1.single_conv.0.moving_mean, shape=(128,), dtype=Float32, requires_grad=False), moving_variance=Parameter (name=evo_net.down3.maxpool_conv.1.single_conv.0.moving_variance, shape=(128,), dtype=Float32, requires_grad=False)>
(1): SpectralNormal<
(parametrizations): Conv2d<input_channels=128, output_channels=256, kernel_size=(3, 3), stride=(1, 1), pad_mode=pad, padding=1, dilation=(1, 1), group=1, has_bias=True, weight_init=<mindspore.common.initializer.HeUniform object at 0xfffe742d13a0>, bias_init=<mindspore.common.initializer.Uniform object at 0xfffe742d1760>, format=NCHW>
>
>
(double_conv): SequentialCell<
(0): BatchNorm2d<num_features=128, eps=1e-05, momentum=0.9, gamma=Parameter (name=evo_net.down3.maxpool_conv.1.double_conv.0.gamma, shape=(128,), dtype=Float32, requires_grad=True), beta=Parameter (name=evo_net.down3.maxpool_conv.1.double_conv.0.beta, shape=(128,), dtype=Float32, requires_grad=True), moving_mean=Parameter (name=evo_net.down3.maxpool_conv.1.double_conv.0.moving_mean, shape=(128,), dtype=Float32, requires_grad=False), moving_variance=Parameter (name=evo_net.down3.maxpool_conv.1.double_conv.0.moving_variance, shape=(128,), dtype=Float32, requires_grad=False)>
(1): ReLU<>
(2): SpectralNormal<
(parametrizations): Conv2d<input_channels=128, output_channels=256, kernel_size=(3, 3), stride=(1, 1), pad_mode=pad, padding=1, dilation=(1, 1), group=1, has_bias=True, weight_init=<mindspore.common.initializer.HeUniform object at 0xfffe742d1a90>, bias_init=<mindspore.common.initializer.Uniform object at 0xfffe742d1b50>, format=NCHW>
>
(3): BatchNorm2d<num_features=256, eps=1e-05, momentum=0.9, gamma=Parameter (name=evo_net.down3.maxpool_conv.1.double_conv.3.gamma, shape=(256,), dtype=Float32, requires_grad=True), beta=Parameter (name=evo_net.down3.maxpool_conv.1.double_conv.3.beta, shape=(256,), dtype=Float32, requires_grad=True), moving_mean=Parameter (name=evo_net.down3.maxpool_conv.1.double_conv.3.moving_mean, shape=(256,), dtype=Float32, requires_grad=False), moving_variance=Parameter (name=evo_net.down3.maxpool_conv.1.double_conv.3.moving_variance, shape=(256,), dtype=Float32, requires_grad=False)>
(4): ReLU<>
(5): SpectralNormal<
(parametrizations): Conv2d<input_channels=256, output_channels=256, kernel_size=(3, 3), stride=(1, 1), pad_mode=pad, padding=1, dilation=(1, 1), group=1, has_bias=True, weight_init=<mindspore.common.initializer.HeUniform object at 0xfffe742d1a00>, bias_init=<mindspore.common.initializer.Uniform object at 0xfffe742d19d0>, format=NCHW>
>
>
>
>
>
(down4): Down<
(maxpool_conv): SequentialCell<
(0): MaxPool2d<kernel_size=2, stride=2, pad_mode=VALID>
(1): DoubleConv<
(single_conv): SequentialCell<
(0): BatchNorm2d<num_features=256, eps=1e-05, momentum=0.9, gamma=Parameter (name=evo_net.down4.maxpool_conv.1.single_conv.0.gamma, shape=(256,), dtype=Float32, requires_grad=True), beta=Parameter (name=evo_net.down4.maxpool_conv.1.single_conv.0.beta, shape=(256,), dtype=Float32, requires_grad=True), moving_mean=Parameter (name=evo_net.down4.maxpool_conv.1.single_conv.0.moving_mean, shape=(256,), dtype=Float32, requires_grad=False), moving_variance=Parameter (name=evo_net.down4.maxpool_conv.1.single_conv.0.moving_variance, shape=(256,), dtype=Float32, requires_grad=False)>
(1): SpectralNormal<
(parametrizations): Conv2d<input_channels=256, output_channels=256, kernel_size=(3, 3), stride=(1, 1), pad_mode=pad, padding=1, dilation=(1, 1), group=1, has_bias=True, weight_init=<mindspore.common.initializer.HeUniform object at 0xfffe742d1f40>, bias_init=<mindspore.common.initializer.Uniform object at 0xfffe742d1f10>, format=NCHW>
>
>
(double_conv): SequentialCell<
(0): BatchNorm2d<num_features=256, eps=1e-05, momentum=0.9, gamma=Parameter (name=evo_net.down4.maxpool_conv.1.double_conv.0.gamma, shape=(256,), dtype=Float32, requires_grad=True), beta=Parameter (name=evo_net.down4.maxpool_conv.1.double_conv.0.beta, shape=(256,), dtype=Float32, requires_grad=True), moving_mean=Parameter (name=evo_net.down4.maxpool_conv.1.double_conv.0.moving_mean, shape=(256,), dtype=Float32, requires_grad=False), moving_variance=Parameter (name=evo_net.down4.maxpool_conv.1.double_conv.0.moving_variance, shape=(256,), dtype=Float32, requires_grad=False)>
(1): ReLU<>
(2): SpectralNormal<
(parametrizations): Conv2d<input_channels=256, output_channels=256, kernel_size=(3, 3), stride=(1, 1), pad_mode=pad, padding=1, dilation=(1, 1), group=1, has_bias=True, weight_init=<mindspore.common.initializer.HeUniform object at 0xfffe74243130>, bias_init=<mindspore.common.initializer.Uniform object at 0xfffe742431f0>, format=NCHW>
>
(3): BatchNorm2d<num_features=256, eps=1e-05, momentum=0.9, gamma=Parameter (name=evo_net.down4.maxpool_conv.1.double_conv.3.gamma, shape=(256,), dtype=Float32, requires_grad=True), beta=Parameter (name=evo_net.down4.maxpool_conv.1.double_conv.3.beta, shape=(256,), dtype=Float32, requires_grad=True), moving_mean=Parameter (name=evo_net.down4.maxpool_conv.1.double_conv.3.moving_mean, shape=(256,), dtype=Float32, requires_grad=False), moving_variance=Parameter (name=evo_net.down4.maxpool_conv.1.double_conv.3.moving_variance, shape=(256,), dtype=Float32, requires_grad=False)>
(4): ReLU<>
(5): SpectralNormal<
(parametrizations): Conv2d<input_channels=256, output_channels=256, kernel_size=(3, 3), stride=(1, 1), pad_mode=pad, padding=1, dilation=(1, 1), group=1, has_bias=True, weight_init=<mindspore.common.initializer.HeUniform object at 0xfffe742d1940>, bias_init=<mindspore.common.initializer.Uniform object at 0xfffe742d16d0>, format=NCHW>
>
>
>
>
>
(up1): Up<
(up): Upsample<>
(conv): DoubleConv<
(single_conv): SequentialCell<
(0): BatchNorm2d<num_features=512, eps=1e-05, momentum=0.9, gamma=Parameter (name=evo_net.up1.conv.single_conv.0.gamma, shape=(512,), dtype=Float32, requires_grad=True), beta=Parameter (name=evo_net.up1.conv.single_conv.0.beta, shape=(512,), dtype=Float32, requires_grad=True), moving_mean=Parameter (name=evo_net.up1.conv.single_conv.0.moving_mean, shape=(512,), dtype=Float32, requires_grad=False), moving_variance=Parameter (name=evo_net.up1.conv.single_conv.0.moving_variance, shape=(512,), dtype=Float32, requires_grad=False)>
(1): SpectralNormal<
(parametrizations): Conv2d<input_channels=512, output_channels=128, kernel_size=(3, 3), stride=(1, 1), pad_mode=pad, padding=1, dilation=(1, 1), group=1, has_bias=True, weight_init=<mindspore.common.initializer.HeUniform object at 0xfffe74243760>, bias_init=<mindspore.common.initializer.Uniform object at 0xfffe74243790>, format=NCHW>
>
>
(double_conv): SequentialCell<
(0): BatchNorm2d<num_features=512, eps=1e-05, momentum=0.9, gamma=Parameter (name=evo_net.up1.conv.double_conv.0.gamma, shape=(512,), dtype=Float32, requires_grad=True), beta=Parameter (name=evo_net.up1.conv.double_conv.0.beta, shape=(512,), dtype=Float32, requires_grad=True), moving_mean=Parameter (name=evo_net.up1.conv.double_conv.0.moving_mean, shape=(512,), dtype=Float32, requires_grad=False), moving_variance=Parameter (name=evo_net.up1.conv.double_conv.0.moving_variance, shape=(512,), dtype=Float32, requires_grad=False)>
(1): ReLU<>
(2): SpectralNormal<
(parametrizations): Conv2d<input_channels=512, output_channels=256, kernel_size=(3, 3), stride=(1, 1), pad_mode=pad, padding=1, dilation=(1, 1), group=1, has_bias=True, weight_init=<mindspore.common.initializer.HeUniform object at 0xfffe74243670>, bias_init=<mindspore.common.initializer.Uniform object at 0xfffe742434c0>, format=NCHW>
>
(3): BatchNorm2d<num_features=256, eps=1e-05, momentum=0.9, gamma=Parameter (name=evo_net.up1.conv.double_conv.3.gamma, shape=(256,), dtype=Float32, requires_grad=True), beta=Parameter (name=evo_net.up1.conv.double_conv.3.beta, shape=(256,), dtype=Float32, requires_grad=True), moving_mean=Parameter (name=evo_net.up1.conv.double_conv.3.moving_mean, shape=(256,), dtype=Float32, requires_grad=False), moving_variance=Parameter (name=evo_net.up1.conv.double_conv.3.moving_variance, shape=(256,), dtype=Float32, requires_grad=False)>
(4): ReLU<>
(5): SpectralNormal<
(parametrizations): Conv2d<input_channels=256, output_channels=128, kernel_size=(3, 3), stride=(1, 1), pad_mode=pad, padding=1, dilation=(1, 1), group=1, has_bias=True, weight_init=<mindspore.common.initializer.HeUniform object at 0xfffe74243b50>, bias_init=<mindspore.common.initializer.Uniform object at 0xfffe74243c10>, format=NCHW>
>
>
>
>
(up2): Up<
(up): Upsample<>
(conv): DoubleConv<
(single_conv): SequentialCell<
(0): BatchNorm2d<num_features=256, eps=1e-05, momentum=0.9, gamma=Parameter (name=evo_net.up2.conv.single_conv.0.gamma, shape=(256,), dtype=Float32, requires_grad=True), beta=Parameter (name=evo_net.up2.conv.single_conv.0.beta, shape=(256,), dtype=Float32, requires_grad=True), moving_mean=Parameter (name=evo_net.up2.conv.single_conv.0.moving_mean, shape=(256,), dtype=Float32, requires_grad=False), moving_variance=Parameter (name=evo_net.up2.conv.single_conv.0.moving_variance, shape=(256,), dtype=Float32, requires_grad=False)>
(1): SpectralNormal<
(parametrizations): Conv2d<input_channels=256, output_channels=64, kernel_size=(3, 3), stride=(1, 1), pad_mode=pad, padding=1, dilation=(1, 1), group=1, has_bias=True, weight_init=<mindspore.common.initializer.HeUniform object at 0xfffe74243cd0>, bias_init=<mindspore.common.initializer.Uniform object at 0xfffe74243df0>, format=NCHW>
>
>
(double_conv): SequentialCell<
(0): BatchNorm2d<num_features=256, eps=1e-05, momentum=0.9, gamma=Parameter (name=evo_net.up2.conv.double_conv.0.gamma, shape=(256,), dtype=Float32, requires_grad=True), beta=Parameter (name=evo_net.up2.conv.double_conv.0.beta, shape=(256,), dtype=Float32, requires_grad=True), moving_mean=Parameter (name=evo_net.up2.conv.double_conv.0.moving_mean, shape=(256,), dtype=Float32, requires_grad=False), moving_variance=Parameter (name=evo_net.up2.conv.double_conv.0.moving_variance, shape=(256,), dtype=Float32, requires_grad=False)>
(1): ReLU<>
(2): SpectralNormal<
(parametrizations): Conv2d<input_channels=256, output_channels=128, kernel_size=(3, 3), stride=(1, 1), pad_mode=pad, padding=1, dilation=(1, 1), group=1, has_bias=True, weight_init=<mindspore.common.initializer.HeUniform object at 0xfffe742340a0>, bias_init=<mindspore.common.initializer.Uniform object at 0xfffe74243e50>, format=NCHW>
>
(3): BatchNorm2d<num_features=128, eps=1e-05, momentum=0.9, gamma=Parameter (name=evo_net.up2.conv.double_conv.3.gamma, shape=(128,), dtype=Float32, requires_grad=True), beta=Parameter (name=evo_net.up2.conv.double_conv.3.beta, shape=(128,), dtype=Float32, requires_grad=True), moving_mean=Parameter (name=evo_net.up2.conv.double_conv.3.moving_mean, shape=(128,), dtype=Float32, requires_grad=False), moving_variance=Parameter (name=evo_net.up2.conv.double_conv.3.moving_variance, shape=(128,), dtype=Float32, requires_grad=False)>
(4): ReLU<>
(5): SpectralNormal<
(parametrizations): Conv2d<input_channels=128, output_channels=64, kernel_size=(3, 3), stride=(1, 1), pad_mode=pad, padding=1, dilation=(1, 1), group=1, has_bias=True, weight_init=<mindspore.common.initializer.HeUniform object at 0xfffe74234130>, bias_init=<mindspore.common.initializer.Uniform object at 0xfffe74234220>, format=NCHW>
>
>
>
>
(up3): Up<
(up): Upsample<>
(conv): DoubleConv<
(single_conv): SequentialCell<
(0): BatchNorm2d<num_features=128, eps=1e-05, momentum=0.9, gamma=Parameter (name=evo_net.up3.conv.single_conv.0.gamma, shape=(128,), dtype=Float32, requires_grad=True), beta=Parameter (name=evo_net.up3.conv.single_conv.0.beta, shape=(128,), dtype=Float32, requires_grad=True), moving_mean=Parameter (name=evo_net.up3.conv.single_conv.0.moving_mean, shape=(128,), dtype=Float32, requires_grad=False), moving_variance=Parameter (name=evo_net.up3.conv.single_conv.0.moving_variance, shape=(128,), dtype=Float32, requires_grad=False)>
(1): SpectralNormal<
(parametrizations): Conv2d<input_channels=128, output_channels=32, kernel_size=(3, 3), stride=(1, 1), pad_mode=pad, padding=1, dilation=(1, 1), group=1, has_bias=True, weight_init=<mindspore.common.initializer.HeUniform object at 0xfffe74243a90>, bias_init=<mindspore.common.initializer.Uniform object at 0xfffe74234460>, format=NCHW>
>
>
(double_conv): SequentialCell<
(0): BatchNorm2d<num_features=128, eps=1e-05, momentum=0.9, gamma=Parameter (name=evo_net.up3.conv.double_conv.0.gamma, shape=(128,), dtype=Float32, requires_grad=True), beta=Parameter (name=evo_net.up3.conv.double_conv.0.beta, shape=(128,), dtype=Float32, requires_grad=True), moving_mean=Parameter (name=evo_net.up3.conv.double_conv.0.moving_mean, shape=(128,), dtype=Float32, requires_grad=False), moving_variance=Parameter (name=evo_net.up3.conv.double_conv.0.moving_variance, shape=(128,), dtype=Float32, requires_grad=False)>
(1): ReLU<>
(2): SpectralNormal<
(parametrizations): Conv2d<input_channels=128, output_channels=64, kernel_size=(3, 3), stride=(1, 1), pad_mode=pad, padding=1, dilation=(1, 1), group=1, has_bias=True, weight_init=<mindspore.common.initializer.HeUniform object at 0xfffe742346d0>, bias_init=<mindspore.common.initializer.Uniform object at 0xfffe74234790>, format=NCHW>
>
(3): BatchNorm2d<num_features=64, eps=1e-05, momentum=0.9, gamma=Parameter (name=evo_net.up3.conv.double_conv.3.gamma, shape=(64,), dtype=Float32, requires_grad=True), beta=Parameter (name=evo_net.up3.conv.double_conv.3.beta, shape=(64,), dtype=Float32, requires_grad=True), moving_mean=Parameter (name=evo_net.up3.conv.double_conv.3.moving_mean, shape=(64,), dtype=Float32, requires_grad=False), moving_variance=Parameter (name=evo_net.up3.conv.double_conv.3.moving_variance, shape=(64,), dtype=Float32, requires_grad=False)>
(4): ReLU<>
(5): SpectralNormal<
(parametrizations): Conv2d<input_channels=64, output_channels=32, kernel_size=(3, 3), stride=(1, 1), pad_mode=pad, padding=1, dilation=(1, 1), group=1, has_bias=True, weight_init=<mindspore.common.initializer.HeUniform object at 0xfffe742348e0>, bias_init=<mindspore.common.initializer.Uniform object at 0xfffe74234400>, format=NCHW>
>
>
>
>
(up4): Up<
(up): Upsample<>
(conv): DoubleConv<
(single_conv): SequentialCell<
(0): BatchNorm2d<num_features=64, eps=1e-05, momentum=0.9, gamma=Parameter (name=evo_net.up4.conv.single_conv.0.gamma, shape=(64,), dtype=Float32, requires_grad=True), beta=Parameter (name=evo_net.up4.conv.single_conv.0.beta, shape=(64,), dtype=Float32, requires_grad=True), moving_mean=Parameter (name=evo_net.up4.conv.single_conv.0.moving_mean, shape=(64,), dtype=Float32, requires_grad=False), moving_variance=Parameter (name=evo_net.up4.conv.single_conv.0.moving_variance, shape=(64,), dtype=Float32, requires_grad=False)>
(1): SpectralNormal<
(parametrizations): Conv2d<input_channels=64, output_channels=32, kernel_size=(3, 3), stride=(1, 1), pad_mode=pad, padding=1, dilation=(1, 1), group=1, has_bias=True, weight_init=<mindspore.common.initializer.HeUniform object at 0xfffe74234be0>, bias_init=<mindspore.common.initializer.Uniform object at 0xfffe74306d30>, format=NCHW>
>
>
(double_conv): SequentialCell<
(0): BatchNorm2d<num_features=64, eps=1e-05, momentum=0.9, gamma=Parameter (name=evo_net.up4.conv.double_conv.0.gamma, shape=(64,), dtype=Float32, requires_grad=True), beta=Parameter (name=evo_net.up4.conv.double_conv.0.beta, shape=(64,), dtype=Float32, requires_grad=True), moving_mean=Parameter (name=evo_net.up4.conv.double_conv.0.moving_mean, shape=(64,), dtype=Float32, requires_grad=False), moving_variance=Parameter (name=evo_net.up4.conv.double_conv.0.moving_variance, shape=(64,), dtype=Float32, requires_grad=False)>
(1): ReLU<>
(2): SpectralNormal<
(parametrizations): Conv2d<input_channels=64, output_channels=32, kernel_size=(3, 3), stride=(1, 1), pad_mode=pad, padding=1, dilation=(1, 1), group=1, has_bias=True, weight_init=<mindspore.common.initializer.HeUniform object at 0xfffe74234940>, bias_init=<mindspore.common.initializer.Uniform object at 0xfffe74234cd0>, format=NCHW>
>
(3): BatchNorm2d<num_features=32, eps=1e-05, momentum=0.9, gamma=Parameter (name=evo_net.up4.conv.double_conv.3.gamma, shape=(32,), dtype=Float32, requires_grad=True), beta=Parameter (name=evo_net.up4.conv.double_conv.3.beta, shape=(32,), dtype=Float32, requires_grad=True), moving_mean=Parameter (name=evo_net.up4.conv.double_conv.3.moving_mean, shape=(32,), dtype=Float32, requires_grad=False), moving_variance=Parameter (name=evo_net.up4.conv.double_conv.3.moving_variance, shape=(32,), dtype=Float32, requires_grad=False)>
(4): ReLU<>
(5): SpectralNormal<
(parametrizations): Conv2d<input_channels=32, output_channels=32, kernel_size=(3, 3), stride=(1, 1), pad_mode=pad, padding=1, dilation=(1, 1), group=1, has_bias=True, weight_init=<mindspore.common.initializer.HeUniform object at 0xfffe74234d90>, bias_init=<mindspore.common.initializer.Uniform object at 0xfffe741a1100>, format=NCHW>
>
>
>
>
(outc): OutConv<
(conv): Conv2d<input_channels=32, output_channels=20, kernel_size=(1, 1), stride=(1, 1), pad_mode=pad, padding=0, dilation=(1, 1), group=1, has_bias=True, weight_init=<mindspore.common.initializer.HeUniform object at 0xfffe742342b0>, bias_init=<mindspore.common.initializer.Uniform object at 0xfffe74234d00>, format=NCHW>
>
(up1_v): Up<
(up): Upsample<>
(conv): DoubleConv<
(single_conv): SequentialCell<
(0): BatchNorm2d<num_features=512, eps=1e-05, momentum=0.9, gamma=Parameter (name=evo_net.up1_v.conv.single_conv.0.gamma, shape=(512,), dtype=Float32, requires_grad=True), beta=Parameter (name=evo_net.up1_v.conv.single_conv.0.beta, shape=(512,), dtype=Float32, requires_grad=True), moving_mean=Parameter (name=evo_net.up1_v.conv.single_conv.0.moving_mean, shape=(512,), dtype=Float32, requires_grad=False), moving_variance=Parameter (name=evo_net.up1_v.conv.single_conv.0.moving_variance, shape=(512,), dtype=Float32, requires_grad=False)>
(1): SpectralNormal<
(parametrizations): Conv2d<input_channels=512, output_channels=128, kernel_size=(3, 3), stride=(1, 1), pad_mode=pad, padding=1, dilation=(1, 1), group=1, has_bias=True, weight_init=<mindspore.common.initializer.HeUniform object at 0xfffe741a1400>, bias_init=<mindspore.common.initializer.Uniform object at 0xfffe741a1430>, format=NCHW>
>
>
(double_conv): SequentialCell<
(0): BatchNorm2d<num_features=512, eps=1e-05, momentum=0.9, gamma=Parameter (name=evo_net.up1_v.conv.double_conv.0.gamma, shape=(512,), dtype=Float32, requires_grad=True), beta=Parameter (name=evo_net.up1_v.conv.double_conv.0.beta, shape=(512,), dtype=Float32, requires_grad=True), moving_mean=Parameter (name=evo_net.up1_v.conv.double_conv.0.moving_mean, shape=(512,), dtype=Float32, requires_grad=False), moving_variance=Parameter (name=evo_net.up1_v.conv.double_conv.0.moving_variance, shape=(512,), dtype=Float32, requires_grad=False)>
(1): ReLU<>
(2): SpectralNormal<
(parametrizations): Conv2d<input_channels=512, output_channels=256, kernel_size=(3, 3), stride=(1, 1), pad_mode=pad, padding=1, dilation=(1, 1), group=1, has_bias=True, weight_init=<mindspore.common.initializer.HeUniform object at 0xfffe741a1610>, bias_init=<mindspore.common.initializer.Uniform object at 0xfffe741a10a0>, format=NCHW>
>
(3): BatchNorm2d<num_features=256, eps=1e-05, momentum=0.9, gamma=Parameter (name=evo_net.up1_v.conv.double_conv.3.gamma, shape=(256,), dtype=Float32, requires_grad=True), beta=Parameter (name=evo_net.up1_v.conv.double_conv.3.beta, shape=(256,), dtype=Float32, requires_grad=True), moving_mean=Parameter (name=evo_net.up1_v.conv.double_conv.3.moving_mean, shape=(256,), dtype=Float32, requires_grad=False), moving_variance=Parameter (name=evo_net.up1_v.conv.double_conv.3.moving_variance, shape=(256,), dtype=Float32, requires_grad=False)>
(4): ReLU<>
(5): SpectralNormal<
(parametrizations): Conv2d<input_channels=256, output_channels=128, kernel_size=(3, 3), stride=(1, 1), pad_mode=pad, padding=1, dilation=(1, 1), group=1, has_bias=True, weight_init=<mindspore.common.initializer.HeUniform object at 0xfffe741a1820>, bias_init=<mindspore.common.initializer.Uniform object at 0xfffe741a18e0>, format=NCHW>
>
>
>
>
(up2_v): Up<
(up): Upsample<>
(conv): DoubleConv<
(single_conv): SequentialCell<
(0): BatchNorm2d<num_features=256, eps=1e-05, momentum=0.9, gamma=Parameter (name=evo_net.up2_v.conv.single_conv.0.gamma, shape=(256,), dtype=Float32, requires_grad=True), beta=Parameter (name=evo_net.up2_v.conv.single_conv.0.beta, shape=(256,), dtype=Float32, requires_grad=True), moving_mean=Parameter (name=evo_net.up2_v.conv.single_conv.0.moving_mean, shape=(256,), dtype=Float32, requires_grad=False), moving_variance=Parameter (name=evo_net.up2_v.conv.single_conv.0.moving_variance, shape=(256,), dtype=Float32, requires_grad=False)>
(1): SpectralNormal<
(parametrizations): Conv2d<input_channels=256, output_channels=64, kernel_size=(3, 3), stride=(1, 1), pad_mode=pad, padding=1, dilation=(1, 1), group=1, has_bias=True, weight_init=<mindspore.common.initializer.HeUniform object at 0xfffe741a19d0>, bias_init=<mindspore.common.initializer.Uniform object at 0xfffe741a19a0>, format=NCHW>
>
>
(double_conv): SequentialCell<
(0): BatchNorm2d<num_features=256, eps=1e-05, momentum=0.9, gamma=Parameter (name=evo_net.up2_v.conv.double_conv.0.gamma, shape=(256,), dtype=Float32, requires_grad=True), beta=Parameter (name=evo_net.up2_v.conv.double_conv.0.beta, shape=(256,), dtype=Float32, requires_grad=True), moving_mean=Parameter (name=evo_net.up2_v.conv.double_conv.0.moving_mean, shape=(256,), dtype=Float32, requires_grad=False), moving_variance=Parameter (name=evo_net.up2_v.conv.double_conv.0.moving_variance, shape=(256,), dtype=Float32, requires_grad=False)>
(1): ReLU<>
(2): SpectralNormal<
(parametrizations): Conv2d<input_channels=256, output_channels=128, kernel_size=(3, 3), stride=(1, 1), pad_mode=pad, padding=1, dilation=(1, 1), group=1, has_bias=True, weight_init=<mindspore.common.initializer.HeUniform object at 0xfffe741a1d00>, bias_init=<mindspore.common.initializer.Uniform object at 0xfffe741a1af0>, format=NCHW>
>
(3): BatchNorm2d<num_features=128, eps=1e-05, momentum=0.9, gamma=Parameter (name=evo_net.up2_v.conv.double_conv.3.gamma, shape=(128,), dtype=Float32, requires_grad=True), beta=Parameter (name=evo_net.up2_v.conv.double_conv.3.beta, shape=(128,), dtype=Float32, requires_grad=True), moving_mean=Parameter (name=evo_net.up2_v.conv.double_conv.3.moving_mean, shape=(128,), dtype=Float32, requires_grad=False), moving_variance=Parameter (name=evo_net.up2_v.conv.double_conv.3.moving_variance, shape=(128,), dtype=Float32, requires_grad=False)>
(4): ReLU<>
(5): SpectralNormal<
(parametrizations): Conv2d<input_channels=128, output_channels=64, kernel_size=(3, 3), stride=(1, 1), pad_mode=pad, padding=1, dilation=(1, 1), group=1, has_bias=True, weight_init=<mindspore.common.initializer.HeUniform object at 0xfffe741a1e20>, bias_init=<mindspore.common.initializer.Uniform object at 0xfffe741a1ee0>, format=NCHW>
>
>
>
>
(up3_v): Up<
(up): Upsample<>
(conv): DoubleConv<
(single_conv): SequentialCell<
(0): BatchNorm2d<num_features=128, eps=1e-05, momentum=0.9, gamma=Parameter (name=evo_net.up3_v.conv.single_conv.0.gamma, shape=(128,), dtype=Float32, requires_grad=True), beta=Parameter (name=evo_net.up3_v.conv.single_conv.0.beta, shape=(128,), dtype=Float32, requires_grad=True), moving_mean=Parameter (name=evo_net.up3_v.conv.single_conv.0.moving_mean, shape=(128,), dtype=Float32, requires_grad=False), moving_variance=Parameter (name=evo_net.up3_v.conv.single_conv.0.moving_variance, shape=(128,), dtype=Float32, requires_grad=False)>
(1): SpectralNormal<
(parametrizations): Conv2d<input_channels=128, output_channels=32, kernel_size=(3, 3), stride=(1, 1), pad_mode=pad, padding=1, dilation=(1, 1), group=1, has_bias=True, weight_init=<mindspore.common.initializer.HeUniform object at 0xfffe741a1a60>, bias_init=<mindspore.common.initializer.Uniform object at 0xfffe741a1ac0>, format=NCHW>
>
>
(double_conv): SequentialCell<
(0): BatchNorm2d<num_features=128, eps=1e-05, momentum=0.9, gamma=Parameter (name=evo_net.up3_v.conv.double_conv.0.gamma, shape=(128,), dtype=Float32, requires_grad=True), beta=Parameter (name=evo_net.up3_v.conv.double_conv.0.beta, shape=(128,), dtype=Float32, requires_grad=True), moving_mean=Parameter (name=evo_net.up3_v.conv.double_conv.0.moving_mean, shape=(128,), dtype=Float32, requires_grad=False), moving_variance=Parameter (name=evo_net.up3_v.conv.double_conv.0.moving_variance, shape=(128,), dtype=Float32, requires_grad=False)>
(1): ReLU<>
(2): SpectralNormal<
(parametrizations): Conv2d<input_channels=128, output_channels=64, kernel_size=(3, 3), stride=(1, 1), pad_mode=pad, padding=1, dilation=(1, 1), group=1, has_bias=True, weight_init=<mindspore.common.initializer.HeUniform object at 0xfffe74111370>, bias_init=<mindspore.common.initializer.Uniform object at 0xfffe74111430>, format=NCHW>
>
(3): BatchNorm2d<num_features=64, eps=1e-05, momentum=0.9, gamma=Parameter (name=evo_net.up3_v.conv.double_conv.3.gamma, shape=(64,), dtype=Float32, requires_grad=True), beta=Parameter (name=evo_net.up3_v.conv.double_conv.3.beta, shape=(64,), dtype=Float32, requires_grad=True), moving_mean=Parameter (name=evo_net.up3_v.conv.double_conv.3.moving_mean, shape=(64,), dtype=Float32, requires_grad=False), moving_variance=Parameter (name=evo_net.up3_v.conv.double_conv.3.moving_variance, shape=(64,), dtype=Float32, requires_grad=False)>
(4): ReLU<>
(5): SpectralNormal<
(parametrizations): Conv2d<input_channels=64, output_channels=32, kernel_size=(3, 3), stride=(1, 1), pad_mode=pad, padding=1, dilation=(1, 1), group=1, has_bias=True, weight_init=<mindspore.common.initializer.HeUniform object at 0xfffe74111490>, bias_init=<mindspore.common.initializer.Uniform object at 0xfffe74111400>, format=NCHW>
>
>
>
>
(up4_v): Up<
(up): Upsample<>
(conv): DoubleConv<
(single_conv): SequentialCell<
(0): BatchNorm2d<num_features=64, eps=1e-05, momentum=0.9, gamma=Parameter (name=evo_net.up4_v.conv.single_conv.0.gamma, shape=(64,), dtype=Float32, requires_grad=True), beta=Parameter (name=evo_net.up4_v.conv.single_conv.0.beta, shape=(64,), dtype=Float32, requires_grad=True), moving_mean=Parameter (name=evo_net.up4_v.conv.single_conv.0.moving_mean, shape=(64,), dtype=Float32, requires_grad=False), moving_variance=Parameter (name=evo_net.up4_v.conv.single_conv.0.moving_variance, shape=(64,), dtype=Float32, requires_grad=False)>
(1): SpectralNormal<
(parametrizations): Conv2d<input_channels=64, output_channels=32, kernel_size=(3, 3), stride=(1, 1), pad_mode=pad, padding=1, dilation=(1, 1), group=1, has_bias=True, weight_init=<mindspore.common.initializer.HeUniform object at 0xfffe741118e0>, bias_init=<mindspore.common.initializer.Uniform object at 0xfffe74111910>, format=NCHW>
>
>
(double_conv): SequentialCell<
(0): BatchNorm2d<num_features=64, eps=1e-05, momentum=0.9, gamma=Parameter (name=evo_net.up4_v.conv.double_conv.0.gamma, shape=(64,), dtype=Float32, requires_grad=True), beta=Parameter (name=evo_net.up4_v.conv.double_conv.0.beta, shape=(64,), dtype=Float32, requires_grad=True), moving_mean=Parameter (name=evo_net.up4_v.conv.double_conv.0.moving_mean, shape=(64,), dtype=Float32, requires_grad=False), moving_variance=Parameter (name=evo_net.up4_v.conv.double_conv.0.moving_variance, shape=(64,), dtype=Float32, requires_grad=False)>
(1): ReLU<>
(2): SpectralNormal<
(parametrizations): Conv2d<input_channels=64, output_channels=32, kernel_size=(3, 3), stride=(1, 1), pad_mode=pad, padding=1, dilation=(1, 1), group=1, has_bias=True, weight_init=<mindspore.common.initializer.HeUniform object at 0xfffe741117c0>, bias_init=<mindspore.common.initializer.Uniform object at 0xfffe74111a00>, format=NCHW>
>
(3): BatchNorm2d<num_features=32, eps=1e-05, momentum=0.9, gamma=Parameter (name=evo_net.up4_v.conv.double_conv.3.gamma, shape=(32,), dtype=Float32, requires_grad=True), beta=Parameter (name=evo_net.up4_v.conv.double_conv.3.beta, shape=(32,), dtype=Float32, requires_grad=True), moving_mean=Parameter (name=evo_net.up4_v.conv.double_conv.3.moving_mean, shape=(32,), dtype=Float32, requires_grad=False), moving_variance=Parameter (name=evo_net.up4_v.conv.double_conv.3.moving_variance, shape=(32,), dtype=Float32, requires_grad=False)>
(4): ReLU<>
(5): SpectralNormal<
(parametrizations): Conv2d<input_channels=32, output_channels=32, kernel_size=(3, 3), stride=(1, 1), pad_mode=pad, padding=1, dilation=(1, 1), group=1, has_bias=True, weight_init=<mindspore.common.initializer.HeUniform object at 0xfffe74111cd0>, bias_init=<mindspore.common.initializer.Uniform object at 0xfffe74111d90>, format=NCHW>
>
>
>
>
(outc_v): OutConv<
(conv): Conv2d<input_channels=32, output_channels=40, kernel_size=(1, 1), stride=(1, 1), pad_mode=pad, padding=0, dilation=(1, 1), group=1, has_bias=True, weight_init=<mindspore.common.initializer.HeUniform object at 0xfffe74111c10>, bias_init=<mindspore.common.initializer.Uniform object at 0xfffe741116a0>, format=NCHW>
>
>
>
Model Training
[8]:
loss_scale = ms.train.loss_scale_manager.FixedLossScaleManager(loss_scale=2048)
evo_loss_fn = EvolutionLoss(evo_model, config)
trainer = EvolutionTrainer(config, evo_model, evo_loss_fn, logger, loss_scale)
trainer.train()
epoch: 1 step: 398, loss is 6.6358147
epoch: 1 step: 399, loss is 11.47153
epoch: 1 step: 400, loss is 7.2303286
Train epoch time: 1794325.439 ms, per step time: 4485.814 ms
2024-01-31 09:32:00,659 - forecast.py[line:191] - INFO: ================================Start Evaluation================================
2024-01-31 09:32:00,661 - forecast.py[line:192] - INFO: The length of data is: 1
2024-01-31 09:32:04,107 - forecast.py[line:179] - INFO: CSI Neighborhood threshold 16 T+10 min: 0.4054458796087876 T+60 min: 0.16474475307855177 T+120 min: 0.09442198339292594
2024-01-31 09:32:04,181 - forecast.py[line:211] - INFO: ================================End Evaluation================================
epoch: 2 step: 398, loss is 7.7851076
epoch: 2 step: 399, loss is 6.364196
epoch: 2 step: 400, loss is 11.064963
Train epoch time: 1688639.108 ms, per step time: 4221.598 ms
2024-01-31 10:00:12,836 - forecast.py[line:191] - INFO: ================================Start Evaluation================================
2024-01-31 10:00:12,838 - forecast.py[line:192] - INFO: The length of data is: 1
2024-01-31 10:00:16,213 - forecast.py[line:179] - INFO: CSI Neighborhood threshold 16 T+10 min: 0.41709989523909563 T+60 min: 0.16894114336218546 T+120 min: 0.10467792846088278
2024-01-31 10:00:16,291 - forecast.py[line:211] - INFO: ================================End Evaluation================================
epoch: 3 step: 398, loss is 5.818244
epoch: 3 step: 399, loss is 6.7248154
epoch: 3 step: 400, loss is 10.864609
Train epoch time: 1690693.724 ms, per step time: 4226.734 ms
2024-01-31 10:28:27,000 - forecast.py[line:191] - INFO: ================================Start Evaluation================================
2024-01-31 10:28:27,002 - forecast.py[line:192] - INFO: The length of data is: 1
2024-01-31 10:28:30,512 - forecast.py[line:179] - INFO: CSI Neighborhood threshold 16 T+10 min: 0.3993438740760541 T+60 min: 0.16725303802177002 T+120 min: 0.0959103616970071
2024-01-31 10:28:30,583 - forecast.py[line:211] - INFO: ================================End Evaluation================================
epoch: 4 step: 398, loss is 7.879731
epoch: 4 step: 399, loss is 9.270348
epoch: 4 step: 400, loss is 10.59611
Train epoch time: 1687615.214 ms, per step time: 4219.038 ms
2024-01-31 10:56:38,214 - forecast.py[line:191] - INFO: ================================Start Evaluation================================
2024-01-31 10:56:38,216 - forecast.py[line:192] - INFO: The length of data is: 1
2024-01-31 10:56:41,586 - forecast.py[line:179] - INFO: CSI Neighborhood threshold 16 T+10 min: 0.4119070738020246 T+60 min: 0.16328413060990918 T+120 min: 0.10628156308461514
2024-01-31 10:56:41,656 - forecast.py[line:211] - INFO: ================================End Evaluation================================
epoch: 5 step: 398, loss is 14.705731
epoch: 5 step: 399, loss is 10.468576
epoch: 5 step: 400, loss is 6.882686
Train epoch time: 1691109.624 ms, per step time: 4227.774 ms
2024-01-31 11:24:52,789 - forecast.py[line:191] - INFO: ================================Start Evaluation================================
2024-01-31 11:24:52,790 - forecast.py[line:192] - INFO: The length of data is: 1
2024-01-31 11:24:56,256 - forecast.py[line:179] - INFO: CSI Neighborhood threshold 16 T+10 min: 0.4118216017414204 T+60 min: 0.15679251293172677 T+120 min: 0.10144497020636714
2024-01-31 11:24:56,322 - forecast.py[line:211] - INFO: ================================End Evaluation================================
Evaluation and Visualization
After training, we use the checkpoint for inference. The visualization of predictions, ground truth and their error is shown below.
[10]:
config["data"]["batch_size"] = 1
config["summary"]["visual"] = True
params = load_checkpoint('./summary/ckpt/evolution-3_200.ckpt')
evo_model.set_train(False)
load_param_into_net(evo_model, params)
evo_inference = EvolutionPredictor(config, evo_model, logger)
[26]:
data_params = config.get("data")
test_dataset_generator = RadarData(data_params, run_mode='test', module_name="evolution")
test_dataset = NowcastDataset(test_dataset_generator,
module_name="evolution",
distribute=train_params.get('distribute', False),
num_workers=data_params.get('num_workers', 1),
shuffle=False)
test_dataset = test_dataset.create_dataset(data_params.get('batch_size', 1))
# data = next(test_dataset.create_dict_iterator())
steps = 1
for d in test_dataset.create_dict_iterator():
if steps == 6:
data = d
break
steps += 1
inputs = data['inputs']
pred = evo_inference.forecast(inputs)
labels = inputs[:, data_params.get("t_in"):]
plt_idx = [x // data_params.get("time_interval") - 1 for x in data_params.get("key_info_timestep", [10, 60, 120])]
plt_img(field=pred[0].asnumpy(), label=labels[0].asnumpy(), idx=plt_idx, fig_name="./evolution_example.png")
Generation
Model Training
[1]:
config["model"]["module_name"] = 'generation'
config["data"]["batch_size"] = 1
config["summary"]["visual"] = False
config["summary"]["save_checkpoint_epochs"] = 1
train_params = config.get("train")
summary_params = config.get("summary")
g_model = GenerationNet(config)
d_model = TemporalDiscriminator(data_params.get("t_in", 9) + data_params.get("t_out", 20))
g_model.set_train()
d_model.set_train()
g_model = amp.auto_mixed_precision(g_model, amp_level=train_params.get("amp_level", 'O2'))
d_model = amp.auto_mixed_precision(d_model, amp_level=train_params.get("amp_level", 'O2'))
loss_scale = nn.DynamicLossScaleUpdateCell(loss_scale_value=2 ** 12, scale_factor=2, scale_window=1000)
g_loss_fn = GenerateLoss(g_model, d_model)
d_loss_fn = DiscriminatorLoss(g_model, d_model)
trainer = GenerationTrainer(config, g_model, d_model, g_loss_fn, d_loss_fn, logger, loss_scale)
trainer.train()
Evaluation and Visualization
After training, we use the checkpoint for inference. The visualization of predictions, ground truth and their error is shown below.
[2]:
config["summary"]["visual"] = True
config["train"]["load_ckpt"] = True
gen_inference = GenerationPredictor(config, g_model, logger)
[3]:
data_params = config.get("data")
model_params = config.get("model")
test_dataset_generator = RadarData(data_params, run_mode='test', module_name="generation")
test_dataset = NowcastDataset(test_dataset_generator,
module_name="generation",
distribute=train_params.get('distribute', False),
num_workers=data_params.get('num_workers', 1),
shuffle=False)
test_dataset = test_dataset.create_dataset(data_params.get('batch_size', 1))
data = next(test_dataset.create_dict_iterator())
inp, evo_result, labels = data.get("inputs"), data.get("evo"), data.get("labels")
noise_scale = data_params.get("noise_scale", 32)
threshold = summary_params.get("csin_threshold", 16)
batch_size = data_params.get("batch_size", 1)
w_size = data_params.get("w_size", 512)
h_size = data_params.get("h_size", 512)
ngf = model_params.get("ngf", 32)
noise = ms.tensor(ms.numpy.randn((batch_size, ngf, h_size // noise_scale, w_size // noise_scale)), inp.dtype)
pred = gen_inference.generator(inp, evo_result, noise)
plt_img(field=pred[0].asnumpy(), label=labels[0].asnumpy(), idx=plt_idx, fig_name="./generation_example.png", evo=evo_result[0].asnumpy() * 128, plot_evo=True)