NowcastNet: 融入物理机制的生成式短临降水预报模型
概述
NowcastNet是由清华大学龙明盛老师团队开发的一个基于雷达数据的短临降水预报模型。 它提供了0-3h的短临降水预报结果，空间分辨率为1km左右。 该模型主要分为evolution和generation两大模块，其中evolution模块融入了物理机制，给出一个粗糙的预测结果。接着，generation模块在此 基础上生成精细化的结果，从而得到最终的降水预报。模型框架图入下图所示(图片来源于论文 Skilful nowcasting of extreme precipitation with NowcastNet)
NowcastNet
NowcastNet网络框架如下图所示
Evolution network：融入物理机制，以历史的\(x_{-T:0}\)为输入，通过U-Net预测动量\(v\)和残差\(s\)，再经过evolution operator得到预测\(x_{1:T}^{''}\)。形式如下：
\[x_{1:T}^{''} = Evolution(x_{-T:0})\]
Nowcast encoder & decoder：采用Semantic Image Synthesis with Spatially-Adaptive Normalization架构把Evolution network的输出\(x_{1:T}^{''}\)做为conditioning进行GAN训练。
技术路径
MindSpore Earth求解该问题的具体流程如下:
创建数据集
模型构建
损失函数
模型训练
模型评估与可视化
训练和测试所用的数据集可以在: 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)
创建数据集
在dataset路径下，下载训练数据集，验证数据集到./dataset目录，修改Nowcastnet.yaml配置文件中的root_dir。
./dataset中的目录结构如下所示：
├── 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>
>
>
>
模型训练
[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================================
evolution评估和可视化
完成训练后，我们使用ckpt进行推理，下述展示了推理的可视化图片
[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
模型训练
[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()
generation评估和可视化
完成训练后，我们使用ckpt进行推理，下述展示了推理的可视化图片
[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)