CudaLSTM

class neuralhydrology.modelzoo.cudalstm.CudaLSTM(cfg: Config)

Bases: BaseModel

LSTM model class, which relies on PyTorch’s CUDA LSTM class.

This class implements the standard LSTM combined with a model head, as specified in the config. Depending on the embedding settings, static and/or dynamic features may or may not be fed through embedding networks before being concatenated and passed through the LSTM. To control the initial forget gate bias, use the config argument initial_forget_bias. Often it is useful to set this value to a positive value at the start of the model training, to keep the forget gate closed and to facilitate the gradient flow. The CudaLSTM class only supports single-timescale predictions. Use MTSLSTM to train a model and get predictions on multiple temporal resolutions at the same time.

Parameters:

cfg (Config) – The run configuration.

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

Perform a forward pass on the CudaLSTM model.

Parameters:

data (Dict[str, torch.Tensor]) – Dictionary, containing input features as key-value pairs.

Returns:

Model outputs and intermediate states as a dictionary.

  • y_hat: model predictions of shape [batch size, sequence length, number of target variables].

  • h_n: hidden state at the last time step of the sequence of shape [batch size, 1, hidden size].

  • c_n: cell state at the last time step of the sequence of shape [batch size, 1, hidden size].

Return type:

Dict[str, torch.Tensor]

module_parts = ['embedding_net', 'lstm', 'head']