MTSLSTM

class neuralhydrology.modelzoo.mtslstm.MTSLSTM(cfg: Config)

Bases: BaseModel

Multi-Timescale LSTM (MTS-LSTM) from Gauch et al. [1].

An LSTM architecture that allows simultaneous prediction at multiple timescales within one model. There are two flavors of this model: MTS-LTSM and sMTS-LSTM (shared MTS-LSTM). The MTS-LSTM processes inputs at low temporal resolutions up to a point in time. Then, the LSTM splits into one branch for each target timescale. Each branch processes the inputs at its respective timescale. Finally, one prediction head per timescale generates the predictions for that timescale based on the LSTM output. Optionally, one can specify: - a different hidden size for each LSTM branch (use a dict in the hidden_size config argument) - different dynamic input variables for each timescale (use a dict in the dynamic_inputs config argument) - the strategy to transfer states from the initial shared low-resolution LSTM to the per-timescale higher-resolution LSTMs. By default, this is a linear transfer layer, but you can specify ‘identity’ to use an identity operation or ‘None’ to turn off any transfer (via the transfer_mtlstm_states config argument).

The sMTS-LSTM variant has the same overall architecture, but the weights of the per-timescale branches (including the output heads) are shared. Thus, unlike MTS-LSTM, the sMTS-LSTM cannot use per-timescale hidden sizes or dynamic input variables.

Parameters:

cfg (Config) – The run configuration.

References

[1]

Gauch, M., Kratzert, F., Klotz, D., Grey, N., Lin, J., and Hochreiter, S.: Rainfall-Runoff Prediction at Multiple Timescales with a Single Long Short-Term Memory Network, arXiv Preprint, https://arxiv.org/abs/2010.07921, 2020.

forward(data: Dict[str, Tensor]) → Dict[str, Tensor]

Perform a forward pass on the MTS-LSTM model.

Parameters:

data (Dict[str, torch.Tensor]) – Input data for the forward pass. See the documentation overview of all models for details on the dict keys.

Returns:

Model predictions for each target timescale.

Return type:

Dict[str, torch.Tensor]

module_parts = ['lstms', 'transfer_fcs', 'heads']