mindearth.cell.DgmrGenerator

class mindearth.cell.DgmrGenerator(forecast_steps=18, in_channels=1, out_channels=256, conv_type='standard', latent_channels=768, context_channels=384, generation_steps=1)[source]

The Dgmr Generator is based on Conditional_Stack, Latent_Stack, Upsample_Stack and ConvGRU, which contain deep residual block. The details can be found in Skilful precipitation nowcasting using deep generative models of radar.

Parameters

forecast_steps (int) – The steps of forecast frames.
in_channels (int) – The channels of input frame.
out_channels (int) – Shape of the output predictions, generally should be same as the input shape.
conv_type (str) – The convolution type.
latent_channels (int) – Latent channels according to network.
context_channels (int) – Context channels according to network.
generation_steps (int) – Number of generation steps to use in forward pass, in paper is 6 and the best is chosen for the loss, this results in huge amounts of GPU memory though, so less might work better for training.

Inputs:

x (Tensor) - Tensor of shape \((batch\_size, input\_frames, out\_channels, height\_size, width\_size)\).

Outputs:

Tensor, the output of Dgmr Generator.

output (Tensor) - Tensor of shape \((batch\_size, output\_frames, out\_channels, height\_size, width\_size)\).

Supported Platforms:

Ascend GPU

Examples

>>> import numpy as np
>>> import mindspore as ms
>>> from mindspore import ops, Tensor
>>> from mindspore.nn import Cell
>>> from mindearth.cell.dgmr.dgmrnet import DgmrGenerator
>>> input_frames = np.random.rand(1, 4, 1, 256, 256).astype(np.float32)
>>> net = DgmrGenerator(
>>>         forecast_steps = 18,
>>>         in_channels = 1,
>>>         out_channels = 256,
>>>         conv_type = "standard",
>>>         latent_channels = 768,
>>>         context_channels = 384,
>>>         generation_steps = 1
>>>     )
>>> out = net(Tensor(input_frames, ms.float32))
>>> print(out.shape)
(1, 18, 1, 256, 256)