ConvLSTM#

class torchgeo.models.ConvLSTM(input_dim, hidden_dim, kernel_size, num_layers, batch_first=True, bias=True, return_all_layers=False)[source]#

Bases: Module

Convolutional LSTM model.

This model is a sequence-processing model that uses convolutional operations within the LSTM cells. It is particularly useful for spatio-temporal data.

If you use this model in your research, please cite the following paper:

Added in version 0.8.

__init__(input_dim, hidden_dim, kernel_size, num_layers, batch_first=True, bias=True, return_all_layers=False)[source]#

Initializes the ConvLSTM model.

Parameters:

  • input_dim (int) – Number of channels in the input.

  • hidden_dim (int | list[int]) – Number of hidden channels. Can be a single int (for all layers) or a list of ints (one for each layer).

  • kernel_size (int | tuple[int, int] | list[int | tuple[int, int]]) –

    Size of the convolutional kernel. Can be:

    • a single integer (for square kernels)

    • a tuple of two integers (for rectangular kernels)

    • a list of integers or tuples (one for each layer)

  • num_layers (int) – Number of LSTM layers stacked on each other.

  • batch_first (bool) – If True, then the input and output tensors are provided as (b, t, c, h, w).

  • bias (bool) – If True, adds a learnable bias to the output.

  • return_all_layers (bool) – If True, will return the list of computations for all layers.

forward(input_tensor, hidden_state=None)[source]#

Forward pass of the ConvLSTM.

Parameters:

  • input_tensor (Tensor) – A 5-D Tensor of shape (t, b, c, h, w) or (b, t, c, h, w).

  • hidden_state (list[tuple[Tensor, Tensor]] | None) – An optional initial hidden state.

Returns:

A tuple containing layer_output_list and last_state_list.

Return type:

tuple[list[Tensor], list[tuple[Tensor, Tensor]]]