Downstream App: Missing Data

Data Module

Neck Module

class sdofmv2.tasks.missing_data.necks.Norm2d(embed_dim: int)

Bases: Module

Applies Layer Normalization over 4D inputs (channels-last).

This module reshapes the input from (B, C, H, W) to (B, H, W, C), applies LayerNorm, and then reshapes it back to (B, C, H, W).

Parameters:

embed_dim (int) – The number of features in the input (C).

Methods

forward(x)

Apply layer normalization to the input.

forward(x)

Apply layer normalization to the input.

Parameters:

x (torch.Tensor) – Input tensor of shape (B, C, H, W).

Returns:

Output tensor of shape (B, C, H, W) with normalized features.

Return type:

torch.Tensor

class sdofmv2.tasks.missing_data.necks.ConvTransformerTokensToEmbeddingNeck(embed_dim: int, output_embed_dim: int, num_frames: int = 1, Hp: int = 14, Wp: int = 14, drop_cls_token: bool = True)

Bases: Module

Neck that transforms the token-based output of transformer into a single embedding suitable for processing with standard layers. Performs 4 ConvTranspose2d operations on the rearranged input with kernel_size=2 and stride=2

Methods

forward(x, ids_restore)

Transform transformer tokens back to spatial embeddings.

forward(x, ids_restore)

Transform transformer tokens back to spatial embeddings.

Parameters:

• x (torch.Tensor) – Token embeddings of shape (batch, num_tokens, token_dim).

• ids_restore (torch.Tensor) – Indices for reordering tokens to original positions.

Returns:

Upsampled embeddings of shape (batch, out_channels, 16*H, 16*W).

Return type:

torch.Tensor

Wrap Encoder

class sdofmv2.tasks.missing_data.wrap_encoder.WrapEncoder(encoder: Module)

Bases: Module

A wrapper for Prithvi-style encoders to handle temporal dimensions.

This class ensures that 4D input tensors (B, C, H, W) are correctly reshaped into 5D tensors (B, C, T, H, W) before being passed to the encoder. It also manages the extraction of intermediate features.

encoder

The underlying encoder module (e.g., a Prithvi ViT).

Methods

forward(x)	Performs a forward pass through the encoder.
forward_features(x, n[, mask_ratio, ...])	Extracts intermediate features from specific layers of the encoder.

forward(x: Tensor) → Tensor

Performs a forward pass through the encoder.

If the input is 4D, a singleton temporal dimension is added. The temporal dimension is squeezed from the output before returning.

Parameters:

x – Input tensor of shape (B, C, H, W) or (B, C, T, H, W).

Returns:

The encoded features as a torch.Tensor.

forward_features(x: Tensor, n: list[int], mask_ratio: float = 0.0, reshape: bool = True, norm: bool = False) → list[Tensor]

Extracts intermediate features from specific layers of the encoder.

Parameters:

• x – Input tensor of shape (B, C, H, W) or (B, C, T, H, W).

• n – A list of layer indices from which to extract features.

• mask_ratio – The fraction of patches to mask during the forward pass.

• reshape – Whether to reshape the output features into a spatial grid.

• norm – Whether to apply normalization to the extracted features.

Returns:

A list of tensors containing the intermediate features from the requested layers.