minerva.models.ssl ================== .. py:module:: minerva.models.ssl Submodules ---------- .. toctree:: :maxdepth: 1 /autoapi/minerva/models/ssl/autoencoder/index /autoapi/minerva/models/ssl/barlowtwins/index /autoapi/minerva/models/ssl/byol/index /autoapi/minerva/models/ssl/cpc/index /autoapi/minerva/models/ssl/diet/index /autoapi/minerva/models/ssl/fastsiam/index /autoapi/minerva/models/ssl/lfr/index /autoapi/minerva/models/ssl/simclr/index /autoapi/minerva/models/ssl/simsiam/index /autoapi/minerva/models/ssl/tfc/index /autoapi/minerva/models/ssl/tnc/index /autoapi/minerva/models/ssl/topological_autoencoder/index /autoapi/minerva/models/ssl/vitmae/index Classes ------- .. autoapisummary:: minerva.models.ssl.BYOL minerva.models.ssl.LearnFromRandomnessModel minerva.models.ssl.RepeatedModuleList minerva.models.ssl.SimCLR Package Contents ---------------- .. py:class:: BYOL(backbone = None, projection_head = None, prediction_head = None, learning_rate = 0.001, schedule = 90000, criterion = None, optimizer = torch.optim.Adam, optimizer_kwargs = None) Bases: :py:obj:`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 :class:`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 :class:`~minerva.losses.negative_cossine_similatiry.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 ---------- backbone : nn.Module, optional Feature extractor network. Defaults to :class:`~minerva.models.nets.image.deeplabv3.DeepLabV3Backbone`. projection_head : nn.Module, optional Projection head mapping encoder features to latent space. If None, a default 3-layer MLP is used. prediction_head : nn.Module, optional Prediction head mapping projected features to target space. If None, a default 2-layer MLP is used. learning_rate : float, default=1e-3 Learning rate for optimizer. schedule : int, default=90000 Number of training steps over which to apply cosine momentum schedule. criterion : nn.Module, optional Loss function. Defaults to :class:`~minerva.losses.negative_cossine_similatiry.NegativeCosineSimilarity`. optimizer : type, optional Optimizer class. Defaults to :class:`torch.optim.Adam` if not provided. optimizer_kwargs : dict, 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 :class:`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 ---------- backbone : torch.nn.Module or LoadableModule The backbone (feature extractor) model. fc : torch.nn.Module or LoadableModule The fully connected head. Use nn.Identity() if not required. loss_fn : torch.nn.Module Loss function to optimize during training. adapter : Callable, optional Function to transform backbone outputs before feeding into `fc`. learning_rate : float, default=1e-3 Learning rate used for optimization. flatten : bool, default=True If True, flattens backbone outputs before `fc`. train_metrics : dict, optional TorchMetrics dictionary for training evaluation. val_metrics : dict, optional TorchMetrics dictionary for validation evaluation. test_metrics : dict, optional TorchMetrics dictionary for test evaluation. freeze_backbone : bool, 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_kwargs : dict, optional Additional kwargs passed to the optimizer constructor. lr_scheduler : type, 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_kwargs : dict, optional Additional kwargs passed to the scheduler constructor. .. py:method:: _default_prediction_head() Creates the default prediction head used in BYOL. .. py:method:: _default_projection_head() Creates the default projection head used in BYOL. .. py:method:: _loss_func(outputs, targets=None) Calculate the loss between the output and the input data. Parameters ---------- y_hat : torch.Tensor The output data from the forward pass. y : torch.Tensor The input data/label. Returns ------- torch.Tensor The loss value. .. py:attribute:: backbone .. py:attribute:: backbone_momentum .. py:method:: cosine_schedule(step, max_steps, start_value, end_value, period = None) Uses cosine decay to gradually modify `start_value` to reach `end_value`. Parameters ---------- step : int Current step number. max_steps : int Total number of steps. start_value : float Starting value. end_value : float Target value. period : Optional[int] Steps over which cosine decay completes a full cycle. Defaults to max_steps. Returns ------- float Cosine decay value. .. py:attribute:: criterion .. py:method:: deactivate_requires_grad(model) Freezes the weights of the model. Parameters ---------- model : nn.Module Model to freeze. .. py:method:: forward(x) Forward pass for the BYOL model. Parameters ---------- x : Tensor Input image tensor. Returns ------- Tensor Output tensor after passing through the backbone, projection, and prediction heads. .. py:method:: forward_momentum(x) Forward pass using momentum encoder. Parameters ---------- x : Tensor Input image tensor. Returns ------- Tensor Output tensor after passing through the momentum backbone and projection head. .. py:attribute:: prediction_head .. py:attribute:: projection_head .. py:attribute:: projection_head_momentum .. py:attribute:: schedule_length :value: 90000 .. py:method:: training_step(batch, batch_idx) Overrides SimpleSupervisedModel's step for BYOL. .. py:method:: update_momentum(model, model_ema, m) Updates model weights using momentum. Parameters ---------- model : nn.Module Original model. model_ema : nn.Module Momentum model. m : float Momentum factor. .. py:class:: 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) Bases: :py:obj:`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. .. py:method:: _loss_from_targets(y_pred, y_proj) 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_pred : torch.Tensor The predictions tensors. y_proj : torch.Tensor The projections tensors. .. py:method:: _select_targets(sample_data) 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 .. py:method:: _single_step(batch, batch_idx, step_name) Perform a single training/validation/test step, computing and logging the loss. Parameters ---------- batch : torch.Tensor The input batch of data. batch_idx : int The index of the batch. step_name : str The name of the step (train, val, test). Returns ------- torch.Tensor The loss value for the batch. .. py:attribute:: adapter :value: None .. py:attribute:: backbone .. py:attribute:: backbone_training_steps_counter :value: 1 .. py:method:: configure_optimizers() Configure the optimizer for the model. This method sets up the optimizer for the model's parameters, excluding the projectors. .. py:attribute:: flatten :value: False .. py:method:: forward(x) Forward pass through the network. Parameters ---------- x : torch.Tensor The input data. Returns ------- torch.Tensor The predicted output and projected input. .. py:attribute:: freeze_backbone :value: False .. py:attribute:: learning_rate :value: 0.001 .. py:attribute:: loss_fn .. py:attribute:: max_backbone_training_steps :value: None .. py:attribute:: num_targets :value: None .. py:method:: on_train_batch_end(outputs, batch, batch_idx) Updates the backbone training steps counter only if the backbone is not frozen. .. py:method:: on_train_batch_start(batch, batch_idx) 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. .. py:method:: on_train_epoch_start() 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. .. py:attribute:: predictor_training_epochs :value: None .. py:attribute:: predictors .. py:attribute:: projectors .. py:attribute:: selection_batch_size :value: 128 .. py:method:: setup(stage) 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. .. py:method:: training_step(batch, batch_idx) Perform a training step using the '_single_step' method. Parameters ---------- batch : torch.Tensor The input batch of data. batch_idx : int The index of the batch. Returns ------- torch.Tensor The loss value for the batch. .. py:method:: validation_step(batch, batch_idx) Perform a validation step using the '_single_step' method. Parameters ---------- batch : torch.Tensor The input batch of data. batch_idx : int The index of the batch. Returns ------- torch.Tensor The loss value for the batch. .. py:attribute:: weight_decay :value: 0.0 .. py:class:: RepeatedModuleList(size, cls, *args, **kwargs) Bases: :py:obj:`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) ) ) .. py:class:: SimCLR(backbone, projection_head, flatten = True, temperature = 0.5, lr = 0.001) Bases: :py:obj:`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 :meth:`to`, etc. .. note:: As per the example above, an ``__init__()`` call to the parent class must be made before assignment on the child. :ivar training: Boolean represents whether this module is in training or evaluation mode. :vartype training: bool Initializes the SimCLR model. Parameters ---------- backbone : nn.Module Backbone model for feature extraction. projection_head : nn.Module Projection head model. flatten : bool, optional, default=True Whether to flatten the output of the backbone model, by default True temperature : float, optional, default=0.5 Temperature for the NT-Xent loss, by default 0.5 lr : float, optional, default=1e-3 Learning rate for the optimizer, by default 1e-3 .. py:method:: _single_step(batch) Performs a single forward and loss computation step. Parameters ---------- batch : Tuple[Tuple[Tensor, Tensor], Any] Input batch containing images and optional labels. Returns ------- Tensor Computed loss for the batch. .. py:attribute:: backbone .. py:method:: configure_optimizers() Configures the optimizer for training. Returns ------- torch.optim.Optimizer Optimizer instance. .. py:attribute:: flatten :value: True .. py:method:: forward(x) Forward pass through the SimCLR model. Parameters ---------- x : Tuple[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). .. py:attribute:: loss .. py:attribute:: lr :value: 0.001 .. py:method:: predict_step(batch, batch_idx, dataloader_idx = None) Predict step. Parameters ---------- batch : Tuple[Tuple[Tensor, Tensor], Any] Input batch containing images and optional labels. batch_idx : int Index of the current batch. dataloader_idx : Optional[int], optional Index of the dataloader, by default None Returns ------- Tensor Computed loss for the batch. .. py:attribute:: projector .. py:method:: training_step(batch, batch_idx) Training step. Parameters ---------- batch : Tuple[Tuple[Tensor, Tensor], Any] Input batch containing images and optional labels. batch_idx : int Index of the current batch. Returns ------- Tensor Computed loss for the batch. .. py:method:: validation_step(batch, batch_idx) Validation step. Parameters ---------- batch : Tuple[Tuple[Tensor, Tensor], Any] Input batch containing images and optional labels. batch_idx : int Index of the current batch. Returns ------- Tensor Computed loss for the batch.