minerva.models.ssl.byol
Classes
A Bootstrap Your Own Latent (BYOL) model for self-supervised learning. |
|
Prediction head used for BYOL. |
|
Projection head used for BYOL. |
|
Base class for all projection and prediction heads. |
Module Contents
- class minerva.models.ssl.byol.BYOL(backbone=None, learning_rate=0.025, schedule=90000)[source]
Bases:
lightning.LightningModuleA 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
- Parameters:
backbone (Optional[torch.nn.Module])
learning_rate (float)
schedule (int)
- class minerva.models.ssl.byol.BYOLPredictionHead(input_dim=256, hidden_dim=4096, output_dim=256)[source]
Bases:
ProjectionHeadPrediction head used for BYOL. “This MLP consists in a linear layer with output size 4096 followed by batch normalization, rectified linear units (ReLU), and a final linear layer with output dimension 256.” [0] [0]: BYOL, 2020, https://arxiv.org/abs/2006.07733
Initialize internal Module state, shared by both nn.Module and ScriptModule.
- Parameters:
input_dim (int)
hidden_dim (int)
output_dim (int)
- class minerva.models.ssl.byol.BYOLProjectionHead(input_dim=2048, hidden_dim=4096, output_dim=256)[source]
Bases:
ProjectionHeadProjection head used for BYOL. “This MLP consists in a linear layer with output size 4096 followed by batch normalization, rectified linear units (ReLU), and a final linear layer with output dimension 256.” [0] [0]: BYOL, 2020, https://arxiv.org/abs/2006.07733
Initialize internal Module state, shared by both nn.Module and ScriptModule.
- Parameters:
input_dim (int)
hidden_dim (int)
output_dim (int)
- avgpool
- class minerva.models.ssl.byol.ProjectionHead(blocks)[source]
Bases:
torch.nn.ModuleBase class for all projection and prediction heads.
Initialize internal Module state, shared by both nn.Module and ScriptModule.
- Parameters:
blocks (Sequence[Union[Tuple[int, int, Optional[torch.nn.Module], Optional[torch.nn.Module]], Tuple[int, int, Optional[torch.nn.Module], Optional[torch.nn.Module], bool]]])
- layers