minerva.models.ssl¶

Submodules¶

Classes¶

`BYOL`	Bootstrap Your Own Latent (BYOL) model for self-supervised representation learning.
`LearnFromRandomnessModel`	A PyTorch Lightning model for pretraining with the technique 'Learning From Random Projectors'.
`RepeatedModuleList`	A module list with the same module cls, instantiated size times.
`SimCLR`	Base class for all neural network modules.

Package Contents¶

class minerva.models.ssl.BYOL(backbone=None, projection_head=None, prediction_head=None, learning_rate=0.001, schedule=90000, criterion=None, optimizer=torch.optim.Adam, optimizer_kwargs=None)[source]¶

Bases: minerva.models.nets.base.SimpleSupervisedModel

Bootstrap Your Own Latent (BYOL) model for self-supervised representation learning.

This class implements the BYOL framework [1], built on top of SimpleSupervisedModel to reuse its optimizer, logging, and training utilities. Unlike typical supervised models, BYOL does not require labeled data; instead, it learns representations by predicting one augmented view of an image from another, using both an online and a momentum encoder.

The model consists of:

An online encoder: backbone + projection head + prediction head.
A momentum encoder: backbone + projection head (no prediction head), updated using an exponential moving average of the online encoder parameters.

Key features:

Self-supervised loss via NegativeCosineSimilarity
Momentum update schedule using cosine decay.
Default optimizer: Adam with weight_decay=1e-6.
Built-in hooks for momentum update and loss computation.

Parameters¶

backbonenn.Module, optional: Feature extractor network. Defaults to DeepLabV3Backbone.
projection_headnn.Module, optional: Projection head mapping encoder features to latent space. If None, a default 3-layer MLP is used.
prediction_headnn.Module, optional: Prediction head mapping projected features to target space. If None, a default 2-layer MLP is used.
learning_ratefloat, default=1e-3: Learning rate for optimizer.
scheduleint, default=90000: Number of training steps over which to apply cosine momentum schedule.
criterionnn.Module, optional: Loss function. Defaults to NegativeCosineSimilarity.
optimizertype, optional: Optimizer class. Defaults to torch.optim.Adam if not provided.
optimizer_kwargsdict, optional: Extra keyword arguments for the optimizer. By default, uses {"weight_decay": 1e-6}.

Notes¶

Metrics are disabled by default since BYOL is self-supervised.
The fc layer from SimpleSupervisedModel is replaced with nn.Identity() because BYOL uses its own projection/prediction heads.
The forward pass returns predictions from the online encoder; the momentum encoder is used internally for target computation only.

References¶

[1] Grill, J.B., Strub, F., Altché, F., Tallec, C., Richemond, P.H., Buchatskaya, E.,: Doersch, C., Pires, B.A., Guo, Z.D., Azar, M.G., Piot, B., Kavukcuoglu, K., Munos, R., & Valko, M. (2020). Bootstrap Your Own Latent - A New Approach to Self-Supervised Learning. Advances in Neural Information Processing Systems, 33, 21271–21284.

Initializes the supervised model with training components and configs.

Parameters¶

backbonetorch.nn.Module or LoadableModule: The backbone (feature extractor) model.
fctorch.nn.Module or LoadableModule: The fully connected head. Use nn.Identity() if not required.
loss_fntorch.nn.Module: Loss function to optimize during training.
adapterCallable, optional: Function to transform backbone outputs before feeding into fc.
learning_ratefloat, default=1e-3: Learning rate used for optimization.
flattenbool, default=True: If True, flattens backbone outputs before fc.
train_metricsdict, optional: TorchMetrics dictionary for training evaluation.
val_metricsdict, optional: TorchMetrics dictionary for validation evaluation.
test_metricsdict, optional: TorchMetrics dictionary for test evaluation.
freeze_backbonebool, default=False: If True, backbone parameters are frozen during training.
optimizer: type: Optimizer class to be instantiated. By default, it is set to torch.optim.Adam. Should be a subclass of torch.optim.Optimizer (e.g., torch.optim.SGD).
optimizer_kwargsdict, optional: Additional kwargs passed to the optimizer constructor.
lr_schedulertype, optional: Learning rate scheduler class to be instantiated. By default, it is set to None, which means no scheduler will be used. Should be a subclass of torch.optim.lr_scheduler.LRScheduler (e.g., torch.optim.lr_scheduler.StepLR).
lr_scheduler_kwargsdict, optional: Additional kwargs passed to the scheduler constructor.

_default_prediction_head()[source]¶

Creates the default prediction head used in BYOL.

Return type:: torch.nn.Module

_default_projection_head()[source]¶

Creates the default projection head used in BYOL.

Return type:: torch.nn.Module

_loss_func(outputs, targets=None)[source]¶

Calculate the loss between the output and the input data.

Parameters¶

y_hattorch.Tensor: The output data from the forward pass.
ytorch.Tensor: The input data/label.

Returns¶

torch.Tensor: The loss value.

Return type:: torch.Tensor

backbone¶

backbone_momentum¶

cosine_schedule(step, max_steps, start_value, end_value, period=None)[source]¶

Uses cosine decay to gradually modify start_value to reach end_value.

Parameters¶

stepint: Current step number.
max_stepsint: Total number of steps.
start_valuefloat: Starting value.
end_valuefloat: Target value.
periodOptional[int]: Steps over which cosine decay completes a full cycle. Defaults to max_steps.

Returns¶

float: Cosine decay value.

Parameters:

step (int)
max_steps (int)
start_value (float)
end_value (float)
period (Optional[int])

Return type:

float

criterion¶

deactivate_requires_grad(model)[source]¶

Freezes the weights of the model.

Parameters¶

modelnn.Module: Model to freeze.

Parameters:: model (torch.nn.Module)

forward(x)[source]¶

Forward pass for the BYOL model.

Parameters¶

xTensor: Input image tensor.

Returns¶

Tensor: Output tensor after passing through the backbone, projection, and prediction heads.

Parameters:: x (torch.Tensor)
Return type:: torch.Tensor

forward_momentum(x)[source]¶

Forward pass using momentum encoder.

Parameters¶

xTensor: Input image tensor.

Returns¶

Tensor: Output tensor after passing through the momentum backbone and projection head.

Parameters:: x (torch.Tensor)
Return type:: torch.Tensor

prediction_head¶

projection_head¶

projection_head_momentum¶

schedule_length = 90000¶

training_step(batch, batch_idx)[source]¶

Overrides SimpleSupervisedModel’s step for BYOL.

Parameters:

batch (Sequence[torch.Tensor])
batch_idx (int)

Return type:

torch.Tensor

update_momentum(model, model_ema, m)[source]¶

Updates model weights using momentum.

Parameters¶

modelnn.Module: Original model.
model_emann.Module: Momentum model.
mfloat: Momentum factor.

Parameters:

model (torch.nn.Module)
model_ema (torch.nn.Module)
m (float)

Parameters:

backbone (Optional[torch.nn.Module])
projection_head (Optional[torch.nn.Module])
prediction_head (Optional[torch.nn.Module])
learning_rate (float)
schedule (int)
criterion (Optional[torch.nn.Module])
optimizer (type)
optimizer_kwargs (Optional[Dict[str, Any]])

class minerva.models.ssl.LearnFromRandomnessModel(backbone, projectors, predictors, loss_fn=None, num_targets=None, adapter=None, learning_rate=0.001, weight_decay=0.0, flatten=False, predictor_training_epochs=None, max_backbone_training_steps=None, selection_batch_size=128)[source]¶

Bases: lightning.LightningModule

A PyTorch Lightning model for pretraining with the technique ‘Learning From Random Projectors’. When using ‘predictor_training_epochs’, please consider updating your number of training epochs as well. Otherwise, the LFR backbone will be trained for less epochs: - If the total training epochs in your Trainer is 100, and ‘predictor_training_epochs’ is 1, then the

backbone will be trained on the epochs 0, 2, 4, 6, … and 98, resulting in the backbone being effectively trained for 50 epochs instead of the specified 100.

If the total training epochs in your Trainer is 100, and ‘predictor_training_epochs’ is 2, then the backbone will be trained on the epochs 0, 3, 6, 9, … and 99, resulting in the backbone being effectively trained for 34 epochs instead of the specified 100.

In conclusion, consider updating your total number of training epochs to:: Total number of training epochs = (intended backbone training epochs) * (predictor_training_epochs + 1)

References¶

Yi Sui, Tongzi Wu, Jesse C. Cresswell, Ga Wu, George Stein, Xiao Shi Huang, Xiaochen Zhang, Maksims Volkovs. “Self-supervised Representation Learning From Random Data Projectors”, 2024

Initialize the LFR_Model, freezing the projectors. Remember to update your number of training epochs when using ‘predictor_training_epochs’:

Total number of training epochs = (intended backbone training epochs) * (predictor_training_epochs + 1)

Parameters¶

backbone: torch.nn.Module: The backbone neural network for feature extraction.
projectors: torch.nn.ModuleList: A list of projector networks.
predictors: torch.nn.ModuleList: A list of predictor networks.
num_targets: Optional[int]: The number of projectors and predictors to select from the lists provided, using the Fast Determinantal Point Process (DPP) algorithm. All projectors and predictors are used if the value received is None, a negative integer, or an integer greater than the length of the lists.
loss_fn: Optional[torch.nn.Module]: The loss function to optimize, by default None. If None, the BatchWiseBarlowTwinLoss is used.
adapter: Optional[Callable[[torch.Tensor], torch.Tensor]]: An optional adapter network to be used in the model, by default None.
learning_rate: float: The learning rate for the optimizer, by default 1e-3.
weight_decay: float: The weight decay for the optimizer, by default 0.0.
flatten: bool: Whether to flatten the input tensor or not, by default False.
predictor_training_epochs: Optional[int]: The number of epochs to train only the predictors (excluding the backbone), by default None. If None, zero, or negative, both the predictors and backbone are trained in every epoch. If a positive integer is provided, the backbone is trained for one epoch, then frozen, and the predictors are trained alone for the specified number of epochs. This cycle is repeated throughout the training phase.
max_backbone_training_steps: Optional[int]: The number of steps the backbone will be trained, by default None. If None, zero, or negative, no limit is applied. The steps where the backbone is frozen are ignored.
selection_batch_size: int: By default 128. When selecting projectors and predictors, this variable decides how many random samples from the dataset are used in the Fast Determinantal Point Process (DPP) algorithm.

_loss_from_targets(y_pred, y_proj)[source]¶

Computes the average loss between each pair of predictor and projector outputs. This function is isolated from _single_step to make it easier to test independently.

Parameters¶

y_predtorch.Tensor: The predictions tensors.
y_projtorch.Tensor: The projections tensors.

Parameters:

y_pred (torch.Tensor)
y_proj (torch.Tensor)

_select_targets(sample_data)[source]¶: Select projectors and predictors based on ‘num_targets’ using the Fast Determinantal Point Process (DPP) algorithm and some sample data. Code adapted from https://github.com/layer6ai-labs/lfr/blob/main/ssl_models/lfr.py

_single_step(batch, batch_idx, step_name)[source]¶

Perform a single training/validation/test step, computing and logging the loss.

Parameters¶

batchtorch.Tensor: The input batch of data.
batch_idxint: The index of the batch.
step_namestr: The name of the step (train, val, test).

Returns¶

torch.Tensor: The loss value for the batch.

Parameters:

batch (torch.Tensor)
batch_idx (int)
step_name (str)

Return type:

torch.Tensor

adapter = None¶

backbone¶

backbone_training_steps_counter = 1¶

configure_optimizers()[source]¶: Configure the optimizer for the model. This method sets up the optimizer for the model’s parameters, excluding the projectors.

flatten = False¶

forward(x)[source]¶

Forward pass through the network.

Parameters¶

xtorch.Tensor: The input data.

Returns¶

torch.Tensor: The predicted output and projected input.

Parameters:: x (torch.Tensor)
Return type:: torch.Tensor

freeze_backbone = False¶

learning_rate = 0.001¶

loss_fn¶

max_backbone_training_steps = None¶

num_targets = None¶

on_train_batch_end(outputs, batch, batch_idx)[source]¶: Updates the backbone training steps counter only if the backbone is not frozen.

on_train_batch_start(batch, batch_idx)[source]¶: If a training steps limit is set, it checks the training step counter at the start of every training batch. If the counter reached the limit, it returns -1, stopping the training.

on_train_epoch_start()[source]¶: Executed at the start of each training epoch. If the predictor training epochs is valid, this function evaluates the current epoch number and freeze or unfreeze the backbone based on it. If the predictor training epochs is None, zero, or negative, the backbone is always trained. In the first epoch, the backbone is trained. In the subsequent ‘predictor_training_epochs’ epochs, it is frozen.

predictor_training_epochs = None¶

predictors¶

projectors¶

selection_batch_size = 128¶

setup(stage)[source]¶: Setup function. If necessary, it picks projectors and predictors based on ‘num_targets’ using the first 128 elements of the training dataset, as used in https://github.com/layer6ai-labs/lfr/blob/main/scripts/har/run_har_diet.sh.

training_step(batch, batch_idx)[source]¶

Perform a training step using the ‘_single_step’ method.

Parameters¶

batchtorch.Tensor: The input batch of data.
batch_idxint: The index of the batch.

Returns¶

torch.Tensor: The loss value for the batch.

Parameters:

batch (torch.Tensor)
batch_idx (int)

validation_step(batch, batch_idx)[source]¶

Perform a validation step using the ‘_single_step’ method.

Parameters¶

batchtorch.Tensor: The input batch of data.
batch_idxint: The index of the batch.

Returns¶

torch.Tensor: The loss value for the batch.

Parameters:

batch (torch.Tensor)
batch_idx (int)

weight_decay = 0.0¶

Parameters:

backbone (torch.nn.Module)
projectors (torch.nn.ModuleList)
predictors (torch.nn.ModuleList)
loss_fn (Optional[torch.nn.Module])
num_targets (Optional[int])
adapter (Optional[Callable[[torch.Tensor], torch.Tensor]])
learning_rate (float)
weight_decay (float)
flatten (bool)
predictor_training_epochs (Optional[int])
max_backbone_training_steps (Optional[int])
selection_batch_size (int)

class minerva.models.ssl.RepeatedModuleList(size, cls, *args, **kwargs)[source]¶

Bases: torch.nn.ModuleList

A module list with the same module cls, instantiated size times.

Initializes the RepeatedModuleList with multiple instances of a given module class.

Parameters¶

size: int: The number of instances to create.
cls: type: The module class to instantiate. Must be a subclass of torch.nn.Module.
*args:: Positional arguments to pass to the module class constructor.
**kwargs:: Keyword arguments to pass to the module class constructor.

Raises¶

AssertionError:: If cls is not a subclass of torch.nn.Module.

Example¶

>>> class SimpleModule(torch.nn.Module):
>>>     def __init__(self, in_features, out_features):
>>>         super().__init__()
>>>         self.linear = torch.nn.Linear(in_features, out_features)
>>>
>>> repeated_modules = RepeatedModuleList(3, SimpleModule, 10, 5)
>>> print(repeated_modules)
RepeatedModuleList(
    (0): SimpleModule(
        (linear): Linear(in_features=10, out_features=5, bias=True)
    )
    (1): SimpleModule(
        (linear): Linear(in_features=10, out_features=5, bias=True)
    )
    (2): SimpleModule(
        (linear): Linear(in_features=10, out_features=5, bias=True)
    )
)

Parameters:

size (int)
cls (type)

class minerva.models.ssl.SimCLR(backbone, projection_head, flatten=True, temperature=0.5, lr=0.001)[source]¶

Bases: lightning.LightningModule

Base class for all neural network modules.

Your models should also subclass this class.

Modules can also contain other Modules, allowing them to be nested in a tree structure. You can assign the submodules as regular attributes:

import torch.nn as nn
import torch.nn.functional as F


class Model(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

Submodules assigned in this way will be registered, and will also have their parameters converted when you call to(), etc.

Note

As per the example above, an __init__() call to the parent class must be made before assignment on the child.

Variables:

training (bool) – Boolean represents whether this module is in training or evaluation mode.

Parameters:

backbone (torch.nn.Module)
projection_head (torch.nn.Module)
flatten (bool)
temperature (float)
lr (float)

Initializes the SimCLR model.

Parameters¶

backbonenn.Module: Backbone model for feature extraction.
projection_headnn.Module: Projection head model.
flattenbool, optional, default=True: Whether to flatten the output of the backbone model, by default True
temperaturefloat, optional, default=0.5: Temperature for the NT-Xent loss, by default 0.5
lrfloat, optional, default=1e-3: Learning rate for the optimizer, by default 1e-3

_single_step(batch)[source]¶

Performs a single forward and loss computation step.

Parameters¶

batchTuple[Tuple[Tensor, Tensor], Any]: Input batch containing images and optional labels.

Returns¶

Tensor: Computed loss for the batch.

Parameters:: batch (Tuple[Tuple[torch.Tensor, torch.Tensor], Any])
Return type:: torch.Tensor

backbone¶

configure_optimizers()[source]¶

Configures the optimizer for training.

Returns¶

torch.optim.Optimizer: Optimizer instance.

Return type:: torch.optim.Optimizer

flatten = True¶

forward(x)[source]¶

Forward pass through the SimCLR model.

Parameters¶

xTuple[Tensor, Tensor]: Input tensor of features with shape (batch_size, input_dim).

Returns¶

Tensor: Output tensor of projected features with shape (batch_size, output_dim).

Parameters:: x (Tuple[torch.Tensor, torch.Tensor])

loss¶

lr = 0.001¶

predict_step(batch, batch_idx, dataloader_idx=None)[source]¶

Predict step.

Parameters¶

batchTuple[Tuple[Tensor, Tensor], Any]: Input batch containing images and optional labels.
batch_idxint: Index of the current batch.
dataloader_idxOptional[int], optional: Index of the dataloader, by default None

Returns¶

Tensor: Computed loss for the batch.

Parameters:

batch (Tuple[Tuple[torch.Tensor, torch.Tensor], Any])
batch_idx (int)
dataloader_idx (Optional[int])

projector¶

training_step(batch, batch_idx)[source]¶

Training step.

Parameters¶

batchTuple[Tuple[Tensor, Tensor], Any]: Input batch containing images and optional labels.
batch_idxint: Index of the current batch.

Returns¶

Tensor: Computed loss for the batch.

Parameters:

batch (Tuple[Tuple[torch.Tensor, torch.Tensor], Any])
batch_idx (int)

Return type:

torch.Tensor

validation_step(batch, batch_idx)[source]¶

Validation step.

Parameters¶

batchTuple[Tuple[Tensor, Tensor], Any]: Input batch containing images and optional labels.
batch_idxint: Index of the current batch.

Returns¶

Tensor: Computed loss for the batch.

Parameters:

batch (Tuple[Tuple[torch.Tensor, torch.Tensor], Any])
batch_idx (int)

Return type:

torch.Tensor