Long Short-Term Memory (LSTM)#

API#

class warp_nn.modules.layers.LSTM(
input_size: int,
hidden_size: int,
*,
num_layers: int = 1,
bidirectional: bool = False,
bias: bool = True,
initialize_parameters: bool = True,
requires_grad: bool = True,
)[source]#

Bases: Module

Apply a multi-layer Long Short-Term Memory (LSTM) to an input sequence.

For each element in the input sequence, each layer computes

\[\begin{split}\begin{array}{ll} i_t = \sigma(W_{ii} \, x_t + b_{ii} + W_{hi} \, h_{t-1} + b_{hi}) \\ f_t = \sigma(W_{if} \, x_t + b_{if} + W_{hf} \, h_{t-1} + b_{hf}) \\ g_t = \tanh(W_{ig} \, x_t + b_{ig} + W_{hg} \, h_{t-1} + b_{hg}) \\ o_t = \sigma(W_{io} \, x_t + b_{io} + W_{ho} \, h_{t-1} + b_{ho}) \\ c_t = f_t \odot c_{t-1} + i_t \odot g_t \\ h_t = o_t \odot \tanh(c_t) \end{array}\end{split}\]

where \(h_t\) and \(c_t\) are the hidden and cell states at time \(t\), \(x_t\) is the input at time \(t\) (the output of the previous layer at time \(t\) for the second and subsequent layers), \(h_{t-1}\) and \(c_{t-1}\) are the hidden and cell states of the layer at time \(t-1\) (or the initial hidden and cell states at time \(0\)), \(i_t\), \(f_t\), \(g_t\) and \(o_t\) are the input, forget, cell and output gates, \(\sigma\) is the sigmoid function and \(\odot\) is the element-wise product.

If bidirectional is true, each layer also processes the sequence in reverse order (from the last element to the first one) with its own parameters, and the hidden states of both directions are concatenated (forward direction first) to form the output of the layer.

Learnable parameters (where k is the layer index, and \(D = 2\) if bidirectional is true, otherwise \(D = 1\)):

Name

Shape

Description

\(W_{ii}, W_{if}, W_{ig}, W_{io}\)

weight_ih_l{k}

(4 * hidden_size, input_size) for k = 0, otherwise (4 * hidden_size, D * hidden_size)

Input-to-hidden weights of the k-th layer

\(W_{hi}, W_{hf}, W_{hg}, W_{ho}\)

weight_hh_l{k}

(4 * hidden_size, hidden_size)

Hidden-to-hidden weights of the k-th layer

\(b_{ii}, b_{if}, b_{ig}, b_{io}\)

bias_ih_l{k}

(4 * hidden_size, 1)

Input-to-hidden bias of the k-th layer. Only if bias is true

\(b_{hi}, b_{hf}, b_{hg}, b_{ho}\)

bias_hh_l{k}

(4 * hidden_size, 1)

Hidden-to-hidden bias of the k-th layer. Only if bias is true

The parameters of the reverse direction (only if bidirectional is true) are named the same, with the _reverse suffix (e.g. weight_ih_l{k}_reverse).

The parameters are initialized from the uniform distribution \(u(-k, k)\) where \(k = \frac{1}{\sqrt{\text{hidden\_size}}}\).

Note

On CUDA, the backward pass (i.e. the gradient computation) of the default (32, 32) tile shape requires more shared memory than available on some devices. In that case, create the module with a smaller tile shape, e.g. within a warp_nn.utils.kernel_config(tile_2d=(16, 16)) context.

Parameters:

  • input_size – The number of input features.

  • hidden_size – The number of hidden features.

  • num_layers – The number of stacked recurrent layers.

  • bidirectional – Whether each layer processes the sequence in both directions.

  • bias – Whether to include a bias term.

  • initialize_parameters – Whether to initialize the parameters with their default/initial values. If false, the parameters are left as uninitialized memory, and reading them (e.g. during a forward pass) before loading their values (e.g. from a state dictionary) is undefined behavior.

  • requires_grad – Whether the parameters and the cached output arrays of the module require gradients.

Raises:

ValueError – If the number of layers is less than 1.

__call__(
input: array,
hidden: tuple[array, array] | None = None,
) → tuple[array, tuple[array, array]][source]#

Forward pass of the module.

Parameters:

  • input – The input sequence array, with shape (batch_size, sequence_length, input_size).

  • hidden – A tuple of the initial hidden state and cell state arrays, both with shapes (D * num_layers, batch_size, hidden_size). If not given, they default to zeros.

Returns:

A tuple of the output sequence array (the hidden states of the last layer for each element in the sequence), with shape (batch_size, sequence_length, D * hidden_size), and a tuple of the final hidden state and cell state arrays (of each layer and direction), both with shapes (D * num_layers, batch_size, hidden_size).

Raises:

ValueError – If the shape of the input, hidden state or cell state array is not valid, or if the hidden and cell states are not given as a tuple.