vp_suite.models.precipitation_nowcasting.ef_blocks

This module contains the building blocks for the video prediction models implemented in https://github.com/Hzzone/Precipitation-Nowcasting (Encoder-Forecaster-based prediction models based on ConvLSTM/TrajGRU).

class Encoder(subnets, rnns)

Bases: torch.nn.modules.module.Module

The Encoder component of the Encoder-Forecaster architecture wraps a recurrent model block to encode the input frames into a state that can then be used for prediction.

__init__(subnets, rnns)

Initializes internal Module state, shared by both nn.Module and ScriptModule.

forward(input)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

forward_by_stage(input, subnet, rnn)
training: bool
class Encoder_Forecaster(device, **model_kwargs)

Bases: vp_suite.base.base_model.VPModel

This is a reimplementation of the Encoder-Forecaster structure, as introduced in “Deep Learning for Precipitation Nowcasting: A Benchmark and A New Model” by Shi et al. (https://arxiv.org/abs/1706.03458). This implementation is based on the PyTorch implementation on https://github.com/Hzzone/Precipitation-Nowcasting .

The Encoder-Forecaster Network stacks multiple convolutional/up-/downsampling and recurrent layers that operate on different spatial scales. This model is an abstract one, and the concrete subclasses will fill the encoder/forecaster with actual model components.

NAME = 'Encoder-Forecaster Structure (Shi et al.)'

The model’s name.

__init__(device, **model_kwargs)

Initializes the model by first setting all model hyperparameters, attributes and the like. Then, the model-specific init will actually create the model from the given hyperparameters

Parameters

  • device (str) – The device identifier for the module.

  • **model_kwargs (Any) – Model arguments such as hyperparameters, input shapes etc.

forward(x, pred_frames=1, **kwargs)

Given an input sequence of t frames, predicts pred_frames (p) frames into the future.

Parameters

  • x (torch.Tensor) – A batch of b sequences of t input frames as a tensor of shape [b, t, c, h, w].

  • pred_frames (int) – The number of frames to predict into the future.

  • () (**kwargs) –

Returns: A batch of sequences of p predicted frames as a tensor of shape [b, p, c, h, w].

pred_1(x, **kwargs)

Given an input sequence of t frames, predicts one single frame into the future.

Parameters

  • x (torch.Tensor) – A batch of b sequences of t input frames as a tensor of shape [b, t, c, h, w].

  • **kwargs (Any) – Optional input parameters such as actions.

Returns: A single frame as a tensor of shape [b, c, h, w].

training: bool
class Forecaster(subnets, rnns)

Bases: torch.nn.modules.module.Module

The Decoder component of the Encoder-Forecaster architecture wraps a recurrent model block to decode the input state into a specified number of predicted future frames.

__init__(subnets, rnns)

Initializes internal Module state, shared by both nn.Module and ScriptModule.

forward(hidden_states, pred_frames)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

forward_by_stage(input, state, pred_frames, subnet, rnn)
training: bool