minerva.models.ssl

Submodules

Classes

`BYOL`	A Bootstrap Your Own Latent (BYOL) model for self-supervised learning.
`SimCLR`	Base class for all neural network modules.

Package Contents

class minerva.models.ssl.BYOL(backbone=None, learning_rate=0.025, schedule=90000)[source]

Bases: lightning.LightningModule

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

References

Grill, J., Strub, F., Altché, F., Tallec, C., Richemond, P. H., Buchatskaya, E., … & 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 BYOL model.

Parameters

backbone: Optional[nn.Module]: The backbone network for feature extraction. Defaults to ResNet18.
learning_rate: float: The learning rate for the optimizer. Defaults to 0.025.
schedule: int: The total number of steps for cosine decay scheduling. Defaults to 90000.

backbone

backbone_momentum

configure_optimizers()[source]

Configures the optimizer for the BYOL model.

Returns

torch.optim.SGD: Optimizer with configured learning rate.

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

learning_rate = 0.025

prediction_head

projection_head

projection_head_momentum

schedule_length = 90000

training_step(batch, batch_idx)[source]

Performs a training step for BYOL.

Parameters

batchSequence[Tensor]: A batch of input pairs (x0, x1).
batch_idxint: Batch index.

Returns

Tensor: The computed loss for the current batch.

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])
learning_rate (float)
schedule (int)

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