Head layer

class hy2dl.modelzoo.head.MDN(cfg: Config)

Bases: Module

Mixture Density Network (MDN) head layer.

This class implements a MDN head layer [1], to map the hidden states produced by the LSTM into the parameters of a mixture distribution.

Parameters:

cfg (Config) – Configuration object containing model hyperparameters and settings.

References

[1]

Klotz, D., Kratzert, F., Gauch, M., Keefe Sampson, A., Brandstetter, J., Klambauer, G., Hochreiter, S., and Nearing, G. (2022). Uncertainty estimation with deep learning for rainfall–runoff modeling. Hydrology and Earth System Sciences, 26(6), 1673–1693. https://doi.org/10.5194/hess-26-1673-2022

forward(x) → dict

Forward pass through the MDN head.

Parameters:

x (torch.Tensor (..., hidden_size))

class hy2dl.modelzoo.head.Regression(cfg: Config)

Bases: Module

Regression head layer.

This class implements a regression head, which maps the hidden states produced by the LSTM into predictions.

Parameters:

cfg (Config) – Configuration object containing model hyperparameters and settings.

forward(x) → dict

Forward pass through the regression head.

Parameters:

x (torch.Tensor (..., hidden_size))

Returns:

Dictionary containing the output tensor and hidden states.

Return type:

dict