biapy.models.memory_bank

This module implements a Memory Bank for contrastive learning, designed to store and manage queues of pixel-level and segment-level features.

The MemoryBank class facilitates the use of a feature queue, which is a common component in self-supervised contrastive learning methods. It allows for the dynamic updating of stored features, ensuring that the bank always contains a diverse and up-to-date set of representations for each class. This is crucial for contrastive losses that rely on a large number of negative samples.

class biapy.models.memory_bank.MemoryBank(num_classes: int = 2, memory_size: int = 5000, feature_dims: int = 256, network_stride: Tuple[int, ...] = (16, 16), pixel_update_freq: int = 10, device: device = device(type='cpu'), ignore_index: int = -1)

Bases: Module

Memory Bank for storing pixel and segment features. Used in contrastive learning to maintain a queue of features.

Parameters:

• num_classes (int) – Number of classes in the dataset.

• memory_size (int) – Size of the memory bank for each class.

• feature_dims (int) – Dimension of the feature vectors stored in the memory bank.

• network_stride (int) – Stride of the network, used to downsample the features.

• pixel_update_freq (int) – Frequency at which pixel features are updated in the memory bank.

• device (torch.device) – Device on which the memory bank is stored (CPU or GPU).

• ignore_index (int, optional) – Value to ignore in the loss calculation. If not provided, no value will be ignored.

dequeue_and_enqueue(keys: Tensor, labels: Tensor)

Dequeue and enqueue features into the memory bank.

Parameters:

• keys (torch.Tensor) – Features to be enqueued, shape (batch_size, classes, H, W) or (batch_size, classes, D, H, W). E.g. (8, 19, 128, 256) for a batch size of 2, 19 classes, and a spatial size of 128x256.

• labels (torch.Tensor) – Ground truth labels, shape (batch_size, 1, H, W) or (batch_size, 1, D, H, W). E.g. (8, 1, 128, 256) for a batch size of 2 and a spatial size of 128x256.