ConvLSTM#

class torchgeo.models.ConvLSTM(input_dim, hidden_dim=64, kernel_size=3, num_layers=1, bias=True, return_all_layers=False, num_classes=1, head_kernel_size=1)[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=64, kernel_size=3, num_layers=1, bias=True, return_all_layers=False, num_classes=1, head_kernel_size=1)[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.

  • 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.

  • num_classes (int) – Optional number of segmentation classes for an attached prediction head.

  • head_kernel_size (int) – Kernel size for the optional segmentation head.

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

Forward pass of ConvLSTM feature extraction.

Added in version 0.10.

Parameters:

  • input_tensor (Tensor) – A 5-D Tensor of shape (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]]]

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

Forward pass for segmentation with the prediction head.

Parameters:

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

  • lengths (Tensor | None) – Optional sequence lengths (B,) before padding/truncation. Values larger than the available sequence length use the final timestep.

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

Returns:

Output tensor of shape (B, num_classes, H, W).

Return type:

Tensor