minerva.models.ssl.byol
Classes
Bootstrap Your Own Latent (BYOL) model for self-supervised learning. |
Module Contents
- class minerva.models.ssl.byol.BYOL(backbone=None, projection_head=None, prediction_head=None, learning_rate=0.001, schedule=90000, criterion=None)[source]
Bases:
lightning.LightningModuleBootstrap 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
- backboneOptional[nn.Module]
The backbone network for feature extraction. Defaults to DeepLabV3Backbone.
- projection_headOptional[nn.Module]
Optional custom projection head module. If None, a default MLP-based projection head is used.
- prediction_headOptional[nn.Module]
Optional custom prediction head module. If None, a default MLP-based prediction head is used.
- learning_ratefloat
The learning rate for the optimizer. Defaults to 1e-3.
- scheduleint
The total number of steps for cosine decay scheduling. Defaults to 90000.
- criterionOptional[Optimizer]
Loss function to use. Defaults to NegativeCosineSimilarity.
- _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
- 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.001
- 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
- Parameters:
backbone (Optional[torch.nn.Module])
projection_head (Optional[torch.nn.Module])
prediction_head (Optional[torch.nn.Module])
learning_rate (Optional[float])
schedule (Optional[int])
criterion (Optional[torch.optim.Optimizer])