mdlearn.nn.modules.lstm_net

LSTMNet module.

Classes

LSTMNet(*args, **kwargs)

class mdlearn.nn.modules.lstm_net.LSTMNet(*args: Any, **kwargs: Any)
__init__(input_dim: int, neurons: list[int] = [32], bias: bool = True, dropout: float = 0.0)

LSTMNet module for easy StackedLSTM network creation. The returned tensor from the forward function, is the hidden state of the final LSTM layer.

Parameters:

  • input_dim (int) – Dimension D of input tensor (N, D) where N is the length of the sequence and D is the dimension of each example.

  • neurons (List[int], default=[32]) – LSTM layers hidden_size.

  • bias (bool, default=True) – If False, then each layer does not use bias weights b_ih and b_hh.

  • dropout (float, default=0.0) – If non-zero, introduces a Dropout layer on the outputs of each LSTM layer except the last layer, with dropout probability equal to dropout.

Note

Bidirectional LSTMs are not currently supported in this module.

Raises:

ValueError – neurons should specify atleast one layer.

forward(x: torch.Tensor) → torch.Tensor

Forward pass through LSTM network.

Parameters:

x (torch.Tensor) – Input data of shape (B, N, D) where B is the batch size, N is the length of the sequence, and D is the dimension.

Returns:

torch.Tensor – The output of the neural network with dimension (batch size, last neuron size).