From 4da08114b61e6042856f5123640faec6f3e58bfd Mon Sep 17 00:00:00 2001 From: gabrielfruet Date: Tue, 11 Aug 2026 18:03:58 -0300 Subject: [PATCH 1/9] Remove deprecated collate functions Collate functions have warned since v1.4.0. Transforms plus MultiViewCollate replace them. lightly-ssl-train now builds a SimCLRTransform from the collate: config namespace instead of an ImageCollateFunction. Every key maps onto a SimCLRTransform argument and all of them are pinned in config.yaml, so the augmentations are unchanged. The namespace keeps its name because lightly-embed reads collate.input_size. plot_augmented_images takes a MultiViewTransform. The view transforms wrap a T.Compose rather than subclassing it, so apply_transform_without_normalize needs to unwrap them, otherwise it returns normalized tensors instead of PIL images. Co-Authored-By: Claude Opus 5 --- docs/source/getting_started/advanced.rst | 12 +- .../code_examples/plot_image_augmentations.py | 16 +- docs/source/lightly.data.rst | 10 +- lightly/__init__.py | 3 +- lightly/cli/config/config.yaml | 2 +- lightly/cli/train_cli.py | 8 +- lightly/data/__init__.py | 15 +- lightly/data/collate.py | 1370 ----------------- lightly/utils/debug.py | 68 +- pyproject.toml | 2 - tests/data/test_data_collate.py | 239 --- tests/data/test_ijepa_collate.py | 15 +- tests/utils/test_debug.py | 48 +- 13 files changed, 92 insertions(+), 1716 deletions(-) delete mode 100644 lightly/data/collate.py delete mode 100644 tests/data/test_data_collate.py diff --git a/docs/source/getting_started/advanced.rst b/docs/source/getting_started/advanced.rst index 9549d7343..7eb7032b2 100644 --- a/docs/source/getting_started/advanced.rst +++ b/docs/source/getting_started/advanced.rst @@ -165,10 +165,6 @@ There are three ways how you can customize augmentations in Lightly\ **SSL**: Previewing Augmentations ^^^^^^^^^^^^^^^^^^^^^^^^ -.. note:: - This section is outdated and still uses the old collate functions which are deprecated - since v1.4.0. We will update this section soon. - It often can be very useful to understand how the image augmentations we pick affect the input dataset. We provide a few helper methods that make it very easy to preview augmentations using Lightly\ **SSL**. @@ -183,7 +179,7 @@ well as their augmentations next to them. :align: center :alt: SimCLR augmentations example - Example augmentations of the `SimCLRCollateFunction` function on images + Example augmentations of the `SimCLRTransform` transform on images from the clothing dataset. The images seem rather blurry! However, we don't want our model to ignore small @@ -194,17 +190,17 @@ details. Let's disable Gaussian Blur and check again: :align: center :alt: SimCLR augmentations example - Example augmentations of the `SimCLRCollateFunction` function on images + Example augmentations of the `SimCLRTransform` transform on images from the clothing dataset. -We can also repeat the experiment for the `DINOCollateFunction` to see what +We can also repeat the experiment for the `DINOTransform` to see what our DINO model would see during training. .. figure:: images/dino_augmentations.jpg :align: center :alt: DINO augmentations example - Example augmentations of the `DINOCollateFunction` function on images + Example augmentations of the `DINOTransform` transform on images from the clothing dataset. diff --git a/docs/source/getting_started/code_examples/plot_image_augmentations.py b/docs/source/getting_started/code_examples/plot_image_augmentations.py index e672dab0b..1e29e4cd8 100644 --- a/docs/source/getting_started/code_examples/plot_image_augmentations.py +++ b/docs/source/getting_started/code_examples/plot_image_augmentations.py @@ -11,16 +11,16 @@ # load the first two images using pillow input_images = [Image.open(fname) for fname in fnames[:2]] -# create our colalte function -collate_fn_simclr = lightly.data.SimCLRCollateFunction() +# create our transform +transform_simclr = lightly.transforms.SimCLRTransform() # plot the images -fig = lightly.utils.debug.plot_augmented_images(input_images, collate_fn_simclr) +fig = lightly.utils.debug.plot_augmented_images(input_images, transform_simclr) # let's disable blur -collate_fn_simclr_no_blur = lightly.data.SimCLRCollateFunction() -fig = lightly.utils.debug.plot_augmented_images(input_images, collate_fn_simclr_no_blur) +transform_simclr_no_blur = lightly.transforms.SimCLRTransform(gaussian_blur=0.0) +fig = lightly.utils.debug.plot_augmented_images(input_images, transform_simclr_no_blur) -# we can also use the DINO collate function instead -collate_fn_dino = lightly.data.DINOCollateFunction() -fig = lightly.utils.debug.plot_augmented_images(input_images, collate_fn_dino) +# we can also use the DINO transform instead +transform_dino = lightly.transforms.DINOTransform() +fig = lightly.utils.debug.plot_augmented_images(input_images, transform_dino) diff --git a/docs/source/lightly.data.rst b/docs/source/lightly.data.rst index ff1e698d8..da1a5ded6 100644 --- a/docs/source/lightly.data.rst +++ b/docs/source/lightly.data.rst @@ -18,15 +18,9 @@ lightly.data -------------- .. note:: - ``IJEPAMaskCollator`` used to live in ``lightly.data.collate``. That import path - still works but warns, and is removed in v1.7.0. Use - ``from lightly.data import IJEPAMaskCollator``. + ``IJEPAMaskCollator`` used to live in ``lightly.data.collate``, which has been + removed. Use ``from lightly.data import IJEPAMaskCollator``. .. autoclass:: lightly.data.ijepa_collate.IJEPAMaskCollator :members: :special-members: __call__ - -.collate: ---------- -.. automodule:: lightly.data.collate - :members: diff --git a/lightly/__init__.py b/lightly/__init__.py index 2f1ed8a15..716220de2 100644 --- a/lightly/__init__.py +++ b/lightly/__init__.py @@ -21,8 +21,7 @@ - **data**: The lightly.data module provides a dataset wrapper and collate functions. The - collate functions are in charge of the data augmentations which are crucial for - self-supervised learning. + collate functions combine the views produced by a transform into a batch. - **loss**: diff --git a/lightly/cli/config/config.yaml b/lightly/cli/config/config.yaml index 08ef1a3cf..1a5c6bcad 100644 --- a/lightly/cli/config/config.yaml +++ b/lightly/cli/config/config.yaml @@ -32,7 +32,7 @@ optimizer: lr: 1. # Learning rate of the optimizer. weight_decay: 0.00001 # L2 penalty. -# collate namespace: Passed to lightly.data.ImageCollateFunction. +# collate namespace: Passed to lightly.transforms.SimCLRTransform. collate: input_size: 64 # Size of the input images in pixels. cj_prob: 0.8 # Probability that color jitter is applied. diff --git a/lightly/cli/train_cli.py b/lightly/cli/train_cli.py index 0cee47d5d..9a343a6ee 100644 --- a/lightly/cli/train_cli.py +++ b/lightly/cli/train_cli.py @@ -26,11 +26,12 @@ load_from_state_dict, load_state_dict_from_url, ) -from lightly.data import ImageCollateFunction, LightlyDataset +from lightly.data import LightlyDataset, MultiViewCollate from lightly.embedding import SelfSupervisedEmbedding from lightly.loss import NTXentLoss from lightly.models import ResNetGenerator from lightly.models.batchnorm import get_norm_layer +from lightly.transforms import SimCLRTransform from lightly.utils.hipify import bcolors @@ -108,13 +109,12 @@ def _train_cli(cfg, is_cli_call=True): criterion = NTXentLoss(**cfg["criterion"]) optimizer = torch.optim.SGD(model.parameters(), **cfg["optimizer"]) - dataset = LightlyDataset(input_dir) + dataset = LightlyDataset(input_dir, transform=SimCLRTransform(**cfg["collate"])) cfg["loader"]["batch_size"] = min(cfg["loader"]["batch_size"], len(dataset)) - collate_fn = ImageCollateFunction(**cfg["collate"]) dataloader = torch.utils.data.DataLoader( - dataset, **cfg["loader"], collate_fn=collate_fn + dataset, **cfg["loader"], collate_fn=MultiViewCollate() ) encoder = SelfSupervisedEmbedding(model, criterion, optimizer, dataloader) diff --git a/lightly/data/__init__.py b/lightly/data/__init__.py index 257626865..119081cb8 100644 --- a/lightly/data/__init__.py +++ b/lightly/data/__init__.py @@ -9,19 +9,6 @@ UnseekableTimestampError, VideoError, ) -from lightly.data.collate import ( - BaseCollateFunction, - DINOCollateFunction, - ImageCollateFunction, - MAECollateFunction, - MoCoCollateFunction, - MSNCollateFunction, - MultiCropCollateFunction, - PIRLCollateFunction, - SimCLRCollateFunction, - SwaVCollateFunction, - VICRegLCollateFunction, - imagenet_normalize, -) from lightly.data.dataset import LightlyDataset from lightly.data.ijepa_collate import IJEPAMaskCollator +from lightly.data.multi_view_collate import MultiViewCollate diff --git a/lightly/data/collate.py b/lightly/data/collate.py deleted file mode 100644 index deeefda66..000000000 --- a/lightly/data/collate.py +++ /dev/null @@ -1,1370 +0,0 @@ -"""Collate Functions""" - -# Copyright (c) 2020. Lightly AG and its affiliates. -# All Rights Reserved - -from typing import TYPE_CHECKING, Any, List, Optional, Tuple, Union -from warnings import warn - -import torch -import torch.nn as nn -import torchvision -from PIL import Image - -from lightly.transforms import GaussianBlur, Jigsaw, RandomSolarization -from lightly.transforms.random_crop_and_flip_with_grid import RandomResizedCropAndFlip -from lightly.transforms.rotation import random_rotation_transform -from lightly.transforms.torchvision_v2_compatibility import torchvision_transforms as T -from lightly.transforms.utils import IMAGENET_NORMALIZE - -imagenet_normalize = IMAGENET_NORMALIZE -# Kept for backwards compatibility - - -class BaseCollateFunction(nn.Module): - """Base class for other collate implementations. - - Takes a batch of images as input and transforms each image into two - different augmentations with the help of random transforms. The images are - then concatenated such that the output batch is exactly twice the length - of the input batch. - - Attributes: - transform: - A set of torchvision transforms which are randomly applied to - each image. - - """ - - def __init__(self, transform: T.Compose): - _deprecation_warning_collate_functions() - super(BaseCollateFunction, self).__init__() - self.transform = transform - - def forward(self, batch: List[Tuple[Image.Image, int, str]]): - """Turns a batch of tuples into a tuple of batches. - - Args: - batch: - A batch of tuples of images, labels, and filenames which - is automatically provided if the dataloader is built from - a LightlyDataset. - - Returns: - A tuple of images, labels, and filenames. The images consist of - two batches corresponding to the two transformations of the - input images. - - Examples: - >>> # define a random transformation and the collate function - >>> transform = ... # some random augmentations - >>> collate_fn = BaseCollateFunction(transform) - >>> - >>> # input is a batch of tuples (here, batch_size = 1) - >>> input = [(img, 0, "my-image.png")] - >>> output = collate_fn(input) - >>> - >>> # output consists of two random transforms of the images, - >>> # the labels, and the filenames in the batch - >>> (img_t0, img_t1), label, filename = output - - """ - batch_size = len(batch) - - # list of transformed images - transforms = [ - self.transform(batch[i % batch_size][0]).unsqueeze_(0) - for i in range(2 * batch_size) - ] - # list of labels - labels = torch.LongTensor([item[1] for item in batch]) - # list of filenames - fnames = [item[2] for item in batch] - - # tuple of transforms - transforms = ( - torch.cat(transforms[:batch_size], 0), - torch.cat(transforms[batch_size:], 0), - ) - - return transforms, labels, fnames - - -class ImageCollateFunction(BaseCollateFunction): - """Implementation of a collate function for images. - - This is an implementation of the BaseCollateFunction with a concrete - set of transforms. - - The set of transforms is inspired by the SimCLR paper as it has shown - to produce powerful embeddings. - - Attributes: - input_size: - Size of the input image in pixels. - cj_prob: - Probability that color jitter is applied. - cj_bright: - How much to jitter brightness. - cj_contrast: - How much to jitter contrast. - cj_sat: - How much to jitter saturation. - cj_hue: - How much to jitter hue. - min_scale: - Minimum size of the randomized crop relative to the input_size. - random_gray_scale: - Probability of conversion to grayscale. - gaussian_blur: - Probability of Gaussian blur. - kernel_size: - Will be deprecated in favor of `sigmas` argument. If set, the old behavior applies and `sigmas` is ignored. - Used to calculate sigma of gaussian blur with kernel_size * input_size. - sigmas: - Tuple of min and max value from which the std of the gaussian kernel is sampled. - Is ignored if `kernel_size` is set. - vf_prob: - Probability that vertical flip is applied. - hf_prob: - Probability that horizontal flip is applied. - rr_prob: - Probability that random rotation is applied. - rr_degrees: - Range of degrees to select from for random rotation. If rr_degrees is None, - images are rotated by 90 degrees. If rr_degrees is a (min, max) tuple, - images are rotated by a random angle in [min, max]. If rr_degrees is a - single number, images are rotated by a random angle in - [-rr_degrees, +rr_degrees]. All rotations are counter-clockwise. - normalize: - Dictionary with 'mean' and 'std' for torchvision.transforms.Normalize. - - """ - - def __init__( - self, - input_size: int = 64, - cj_prob: float = 0.8, - cj_bright: float = 0.7, - cj_contrast: float = 0.7, - cj_sat: float = 0.7, - cj_hue: float = 0.2, - min_scale: float = 0.15, - random_gray_scale: float = 0.2, - gaussian_blur: float = 0.5, - kernel_size: Optional[float] = None, - sigmas: Tuple[float, float] = (0.2, 2), - vf_prob: float = 0.0, - hf_prob: float = 0.5, - rr_prob: float = 0.0, - rr_degrees: Optional[Union[float, Tuple[float, float]]] = None, - normalize: dict = imagenet_normalize, - ): - if isinstance(input_size, tuple): - input_size_ = max(input_size) - else: - input_size_ = input_size - - color_jitter = T.ColorJitter(cj_bright, cj_contrast, cj_sat, cj_hue) - - transform = [ - T.RandomResizedCrop(size=input_size, scale=(min_scale, 1.0)), - T.RandomHorizontalFlip(p=hf_prob), - T.RandomVerticalFlip(p=vf_prob), - random_rotation_transform(rr_prob=rr_prob, rr_degrees=rr_degrees), - T.RandomApply([color_jitter], p=cj_prob), - T.RandomGrayscale(p=random_gray_scale), - GaussianBlur(kernel_size=kernel_size, sigmas=sigmas, prob=gaussian_blur), - T.ToTensor(), - ] - - if normalize: - transform += [T.Normalize(mean=normalize["mean"], std=normalize["std"])] - - transform = T.Compose(transform) - - super(ImageCollateFunction, self).__init__(transform) - - -class MultiViewCollateFunction(nn.Module): - """Generates multiple views for each image in the batch. - - Attributes: - transforms: - List of transformation functions. Each function is used to generate - one view of the back. - - """ - - def __init__(self, transforms: List[T.Compose]): - _deprecation_warning_collate_functions() - super().__init__() - self.transforms = transforms - - def forward(self, batch: List[tuple]): - """Turns a batch of tuples into a tuple of batches. - - Args: - batch: - The input batch. - - Returns: - A (views, labels, fnames) tuple where views is a list of tensors - with each tensor containing one view of the batch. - - """ - views = [] - for transform in self.transforms: - view = torch.stack([transform(img) for img, _, _ in batch]) - views.append(view) - # list of labels - labels = torch.LongTensor([label for _, label, _ in batch]) - # list of filenames - fnames = [fname for _, _, fname in batch] - return views, labels, fnames - - -class SimCLRCollateFunction(ImageCollateFunction): - """Implements the transformations for SimCLR. - - Attributes: - input_size: - Size of the input image in pixels. - cj_prob: - Probability that color jitter is applied. - cj_strength: - Strength of the color jitter. - min_scale: - Minimum size of the randomized crop relative to the input_size. - random_gray_scale: - Probability of conversion to grayscale. - gaussian_blur: - Probability of Gaussian blur. - kernel_size: - Will be deprecated in favor of `sigmas` argument. If set, the old behavior applies and `sigmas` is ignored. - Used to calculate sigma of gaussian blur with kernel_size * input_size. - sigmas: - Tuple of min and max value from which the std of the gaussian kernel is sampled. - Is ignored if `kernel_size` is set. - vf_prob: - Probability that vertical flip is applied. - hf_prob: - Probability that horizontal flip is applied. - rr_prob: - Probability that random rotation is applied. - rr_degrees: - Range of degrees to select from for random rotation. If rr_degrees is None, - images are rotated by 90 degrees. If rr_degrees is a (min, max) tuple, - images are rotated by a random angle in [min, max]. If rr_degrees is a - single number, images are rotated by a random angle in - [-rr_degrees, +rr_degrees]. All rotations are counter-clockwise. - normalize: - Dictionary with 'mean' and 'std' for torchvision.transforms.Normalize. - - Examples: - >>> # SimCLR for ImageNet - >>> collate_fn = SimCLRCollateFunction() - >>> - >>> # SimCLR for CIFAR-10 - >>> collate_fn = SimCLRCollateFunction( - >>> input_size=32, - >>> gaussian_blur=0., - >>> ) - - """ - - def __init__( - self, - input_size: int = 224, - cj_prob: float = 0.8, - cj_strength: float = 0.5, - min_scale: float = 0.08, - random_gray_scale: float = 0.2, - gaussian_blur: float = 0.5, - kernel_size: Optional[float] = None, - sigmas: Tuple[float, float] = (0.2, 2), - vf_prob: float = 0.0, - hf_prob: float = 0.5, - rr_prob: float = 0.0, - rr_degrees: Optional[Union[float, Tuple[float, float]]] = None, - normalize: dict = imagenet_normalize, - ): - super(SimCLRCollateFunction, self).__init__( - input_size=input_size, - cj_prob=cj_prob, - cj_bright=cj_strength * 0.8, - cj_contrast=cj_strength * 0.8, - cj_sat=cj_strength * 0.8, - cj_hue=cj_strength * 0.2, - min_scale=min_scale, - random_gray_scale=random_gray_scale, - gaussian_blur=gaussian_blur, - kernel_size=kernel_size, - sigmas=sigmas, - vf_prob=vf_prob, - hf_prob=hf_prob, - rr_prob=rr_prob, - rr_degrees=rr_degrees, - normalize=normalize, - ) - - -class MoCoCollateFunction(ImageCollateFunction): - """Implements the transformations for MoCo v1. - - For MoCo v2, simply use the SimCLR settings. - - Attributes: - input_size: - Size of the input image in pixels. - cj_prob: - Probability that color jitter is applied. - cj_strength: - Strength of the color jitter. - min_scale: - Minimum size of the randomized crop relative to the input_size. - random_gray_scale: - Probability of conversion to grayscale. - gaussian_blur: - Probability of Gaussian blur. - kernel_size: - Will be deprecated in favor of `sigmas` argument. If set, the old behavior applies and `sigmas` is ignored. - Used to calculate sigma of gaussian blur with kernel_size * input_size. - sigmas: - Tuple of min and max value from which the std of the gaussian kernel is sampled. - Is ignored if `kernel_size` is set. - vf_prob: - Probability that vertical flip is applied. - hf_prob: - Probability that horizontal flip is applied. - rr_prob: - Probability that random rotation is applied. - rr_degrees: - Range of degrees to select from for random rotation. If rr_degrees is None, - images are rotated by 90 degrees. If rr_degrees is a (min, max) tuple, - images are rotated by a random angle in [min, max]. If rr_degrees is a - single number, images are rotated by a random angle in - [-rr_degrees, +rr_degrees]. All rotations are counter-clockwise. - normalize: - Dictionary with 'mean' and 'std' for torchvision.transforms.Normalize. - - Examples: - >>> # MoCo v1 for ImageNet - >>> collate_fn = MoCoCollateFunction() - >>> - >>> # MoCo v1 for CIFAR-10 - >>> collate_fn = MoCoCollateFunction( - >>> input_size=32, - >>> ) - - """ - - def __init__( - self, - input_size: int = 224, - cj_prob: float = 0.8, - cj_strength: float = 0.4, - min_scale: float = 0.2, - random_gray_scale: float = 0.2, - gaussian_blur: float = 0.0, - kernel_size: Optional[float] = None, - sigmas: Tuple[float, float] = (0.2, 2), - vf_prob: float = 0.0, - hf_prob: float = 0.5, - rr_prob: float = 0.0, - rr_degrees: Optional[Union[float, Tuple[float, float]]] = None, - normalize: dict = imagenet_normalize, - ): - super(MoCoCollateFunction, self).__init__( - input_size=input_size, - cj_prob=cj_prob, - cj_bright=cj_strength, - cj_contrast=cj_strength, - cj_sat=cj_strength, - cj_hue=cj_strength, - min_scale=min_scale, - random_gray_scale=random_gray_scale, - gaussian_blur=gaussian_blur, - kernel_size=kernel_size, - sigmas=sigmas, - vf_prob=vf_prob, - hf_prob=hf_prob, - rr_prob=rr_prob, - rr_degrees=rr_degrees, - normalize=normalize, - ) - - -class MultiCropCollateFunction(MultiViewCollateFunction): - """Implements the multi-crop transformations for SwaV. - - Attributes: - crop_sizes: - Size of the input image in pixels for each crop category. - crop_counts: - Number of crops for each crop category. - crop_min_scales: - Min scales for each crop category. - crop_max_scales: - Max_scales for each crop category. - transforms: - Transforms which are applied to all crops. - - """ - - def __init__( - self, - crop_sizes: List[int], - crop_counts: List[int], - crop_min_scales: List[float], - crop_max_scales: List[float], - transforms: T.Compose, - ): - if len(crop_sizes) != len(crop_counts): - raise ValueError( - "Length of crop_sizes and crop_counts must be equal but are" - f" {len(crop_sizes)} and {len(crop_counts)}." - ) - if len(crop_sizes) != len(crop_min_scales): - raise ValueError( - "Length of crop_sizes and crop_min_scales must be equal but are" - f" {len(crop_sizes)} and {len(crop_min_scales)}." - ) - if len(crop_sizes) != len(crop_min_scales): - raise ValueError( - "Length of crop_sizes and crop_max_scales must be equal but are" - f" {len(crop_sizes)} and {len(crop_min_scales)}." - ) - - crop_transforms = [] - for i in range(len(crop_sizes)): - random_resized_crop = T.RandomResizedCrop( - crop_sizes[i], scale=(crop_min_scales[i], crop_max_scales[i]) - ) - - crop_transforms.extend( - [ - T.Compose( - [ - random_resized_crop, - transforms, - ] - ) - ] - * crop_counts[i] - ) - super().__init__(crop_transforms) - - -class SwaVCollateFunction(MultiCropCollateFunction): - """Implements the multi-crop transformations for SwaV. - - Attributes: - crop_sizes: - Size of the input image in pixels for each crop category. - crop_counts: - Number of crops for each crop category. - crop_min_scales: - Min scales for each crop category. - crop_max_scales: - Max_scales for each crop category. - hf_prob: - Probability that horizontal flip is applied. - vf_prob: - Probability that vertical flip is applied. - rr_prob: - Probability that random rotation is applied. - rr_degrees: - Range of degrees to select from for random rotation. If rr_degrees is None, - images are rotated by 90 degrees. If rr_degrees is a (min, max) tuple, - images are rotated by a random angle in [min, max]. If rr_degrees is a - single number, images are rotated by a random angle in - [-rr_degrees, +rr_degrees]. All rotations are counter-clockwise. - cj_prob: - Probability that color jitter is applied. - cj_strength: - Strength of the color jitter. - random_gray_scale: - Probability of conversion to grayscale. - gaussian_blur: - Probability of Gaussian blur. - kernel_size: - Will be deprecated in favor of `sigmas` argument. If set, the old behavior applies and `sigmas` is ignored. - Used to calculate sigma of gaussian blur with kernel_size * input_size. - sigmas: - Tuple of min and max value from which the std of the gaussian kernel is sampled. - Is ignored if `kernel_size` is set. - normalize: - Dictionary with 'mean' and 'std' for torchvision.transforms.Normalize. - - Examples: - >>> # SwaV for Imagenet - >>> collate_fn = SwaVCollateFunction() - >>> - >>> # SwaV w/ 2x160 and 4x96 crops - >>> collate_fn = SwaVCollateFunction( - >>> crop_sizes=[160, 96], - >>> crop_counts=[2, 4], - >>> ) - - """ - - def __init__( - self, - crop_sizes: List[int] = [224, 96], - crop_counts: List[int] = [2, 6], - crop_min_scales: List[float] = [0.14, 0.05], - crop_max_scales: List[float] = [1.0, 0.14], - hf_prob: float = 0.5, - vf_prob: float = 0.0, - rr_prob: float = 0.0, - rr_degrees: Optional[Union[float, Tuple[float, float]]] = None, - cj_prob: float = 0.8, - cj_strength: float = 0.8, - random_gray_scale: float = 0.2, - gaussian_blur: float = 0.5, - kernel_size: Optional[float] = None, - sigmas: Tuple[float, float] = (0.2, 2), - normalize: dict = imagenet_normalize, - ): - color_jitter = T.ColorJitter( - cj_strength, - cj_strength, - cj_strength, - cj_strength / 4.0, - ) - - transforms = T.Compose( - [ - T.RandomHorizontalFlip(p=hf_prob), - T.RandomVerticalFlip(p=vf_prob), - random_rotation_transform(rr_prob=rr_prob, rr_degrees=rr_degrees), - T.ColorJitter(), - T.RandomApply([color_jitter], p=cj_prob), - T.RandomGrayscale(p=random_gray_scale), - GaussianBlur( - kernel_size=kernel_size, sigmas=sigmas, prob=gaussian_blur - ), - T.ToTensor(), - T.Normalize(mean=normalize["mean"], std=normalize["std"]), - ] - ) - - super(SwaVCollateFunction, self).__init__( - crop_sizes=crop_sizes, - crop_counts=crop_counts, - crop_min_scales=crop_min_scales, - crop_max_scales=crop_max_scales, - transforms=transforms, - ) - - -class DINOCollateFunction(MultiViewCollateFunction): - """Implements the global and local view augmentations for DINO [0]. - - This class generates two global and a user defined number of local views - for each image in a batch. The code is adapted from [1]. - - - [0]: DINO, 2021, https://arxiv.org/abs/2104.14294 - - [1]: https://github.com/facebookresearch/dino - - Attributes: - global_crop_size: - Crop size of the global views. - global_crop_scale: - Tuple of min and max scales relative to global_crop_size. - local_crop_size: - Crop size of the local views. - local_crop_scale: - Tuple of min and max scales relative to local_crop_size. - n_local_views: - Number of generated local views. - hf_prob: - Probability that horizontal flip is applied. - vf_prob: - Probability that vertical flip is applied. - rr_prob: - Probability that random rotation is applied. - rr_degrees: - Range of degrees to select from for random rotation. If rr_degrees is None, - images are rotated by 90 degrees. If rr_degrees is a (min, max) tuple, - images are rotated by a random angle in [min, max]. If rr_degrees is a - single number, images are rotated by a random angle in - [-rr_degrees, +rr_degrees]. All rotations are counter-clockwise. - cj_prob: - Probability that color jitter is applied. - cj_bright: - How much to jitter brightness. - cj_contrast: - How much to jitter contrast. - cj_sat: - How much to jitter saturation. - cj_hue: - How much to jitter hue. - random_gray_scale: - Probability of conversion to grayscale. - gaussian_blur: - Tuple of probabilities to apply gaussian blur on the different - views. The input is ordered as follows: - (global_view_0, global_view_1, local_views) - kernel_size: - Will be deprecated in favor of `sigmas` argument. If set, the old behavior applies and `sigmas` is ignored. - Used to calculate sigma of gaussian blur with kernel_size * input_size. - kernel_scale: - Old argument. Value is deprecated in favor of sigmas. If set, the old behavior applies and `sigmas` is ignored. - Used to scale the `kernel_size` of a factor of `kernel_scale` - sigmas: - Tuple of min and max value from which the std of the gaussian kernel is sampled. - Is ignored if `kernel_size` is set. - solarization: - Probability to apply solarization on the second global view. - normalize: - Dictionary with 'mean' and 'std' for torchvision.transforms.Normalize. - - """ - - def __init__( - self, - global_crop_size=224, - global_crop_scale=(0.4, 1.0), - local_crop_size=96, - local_crop_scale=(0.05, 0.4), - n_local_views=6, - hf_prob=0.5, - vf_prob=0, - rr_prob=0, - rr_degrees: Optional[Union[float, Tuple[float, float]]] = None, - cj_prob=0.8, - cj_bright=0.4, - cj_contrast=0.4, - cj_sat=0.2, - cj_hue=0.1, - random_gray_scale=0.2, - gaussian_blur=(1.0, 0.1, 0.5), - kernel_size: Optional[float] = None, - kernel_scale: Optional[float] = None, - sigmas: Tuple[float, float] = (0.1, 2), - solarization_prob=0.2, - normalize=imagenet_normalize, - ): - flip_and_color_jitter = T.Compose( - [ - T.RandomHorizontalFlip(p=hf_prob), - T.RandomVerticalFlip(p=vf_prob), - random_rotation_transform(rr_prob=rr_prob, rr_degrees=rr_degrees), - T.RandomApply( - [ - T.ColorJitter( - brightness=cj_bright, - contrast=cj_contrast, - saturation=cj_sat, - hue=cj_hue, - ) - ], - p=cj_prob, - ), - T.RandomGrayscale(p=random_gray_scale), - ] - ) - normalize = T.Compose( - [ - T.ToTensor(), - T.Normalize(mean=normalize["mean"], std=normalize["std"]), - ] - ) - global_crop = T.RandomResizedCrop( - global_crop_size, - scale=global_crop_scale, - interpolation=Image.BICUBIC, - ) - - # first global crop - global_transform_0 = T.Compose( - [ - global_crop, - flip_and_color_jitter, - GaussianBlur( - kernel_size=kernel_size, - scale=kernel_scale, - sigmas=sigmas, - prob=gaussian_blur[0], - ), - normalize, - ] - ) - - # second global crop - global_transform_1 = T.Compose( - [ - global_crop, - flip_and_color_jitter, - GaussianBlur( - kernel_size=kernel_size, - scale=kernel_scale, - sigmas=sigmas, - prob=gaussian_blur[1], - ), - RandomSolarization(prob=solarization_prob), - normalize, - ] - ) - - # transformation for the local small crops - local_transform = T.Compose( - [ - T.RandomResizedCrop( - local_crop_size, scale=local_crop_scale, interpolation=Image.BICUBIC - ), - flip_and_color_jitter, - GaussianBlur( - kernel_size=kernel_size, - scale=kernel_scale, - sigmas=sigmas, - prob=gaussian_blur[2], - ), - normalize, - ] - ) - local_transforms = [local_transform] * n_local_views - - transforms = [global_transform_0, global_transform_1] - transforms.extend(local_transforms) - super().__init__(transforms) - - -class MAECollateFunction(MultiViewCollateFunction): - """Implements the view augmentation for MAE [0]. - - - [0]: Masked Autoencoder, 2021, https://arxiv.org/abs/2111.06377 - - Attributes: - input_size: - Size of the input image in pixels. - min_scale: - Minimum size of the randomized crop relative to the input_size. - normalize: - Dictionary with 'mean' and 'std' for torchvision.transforms.Normalize. - - """ - - def __init__( - self, - input_size: Union[int, Tuple[int, int]] = 224, - min_scale: float = 0.2, - normalize: dict = imagenet_normalize, - ): - transforms = [ - T.RandomResizedCrop( - input_size, scale=(min_scale, 1.0), interpolation=3 - ), # 3 is bicubic - T.RandomHorizontalFlip(), - T.ToTensor(), - ] - - if normalize: - transforms.append(T.Normalize(mean=normalize["mean"], std=normalize["std"])) - super().__init__([T.Compose(transforms)]) - - def forward(self, batch: List[tuple]): - views, labels, fnames = super().forward(batch) - # Return only first view as MAE needs only a single view per image. - return views[0], labels, fnames - - -class PIRLCollateFunction(nn.Module): - """Implements the transformations for PIRL [0]. The jigsaw augmentation - is applied during the forward pass. - - - [0] PIRL, 2019: https://arxiv.org/abs/1912.01991 - - Attributes: - input_size: - Size of the input image in pixels. - cj_prob: - Probability that color jitter is applied. - cj_bright: - How much to jitter brightness. - cj_contrast: - How much to jitter contrast. - cj_sat: - How much to jitter saturation. - cj_hue: - How much to jitter hue. - min_scale: - Minimum size of the randomized crop relative to the input_size. - random_gray_scale: - Probability of conversion to grayscale. - hf_prob: - Probability that horizontal flip is applied. - n_grid: - Sqrt of the number of grids in the jigsaw image. - normalize: - Dictionary with 'mean' and 'std' for torchvision.transforms.Normalize. - - Examples: - >>> # PIRL for ImageNet - >>> collate_fn = PIRLCollateFunction() - >>> - >>> # PIRL for CIFAR-10 - >>> collate_fn = PIRLCollateFunction( - >>> input_size=32, - >>> ) - - """ - - def __init__( - self, - input_size: int = 64, - cj_prob: float = 0.8, - cj_bright: float = 0.4, - cj_contrast: float = 0.4, - cj_sat: float = 0.4, - cj_hue: float = 0.4, - min_scale: float = 0.08, - random_gray_scale: float = 0.2, - hf_prob: float = 0.5, - n_grid: int = 3, - normalize: dict = imagenet_normalize, - ): - _deprecation_warning_collate_functions() - super(PIRLCollateFunction, self).__init__() - - if isinstance(input_size, tuple): - input_size_ = max(input_size) - else: - input_size_ = input_size - - color_jitter = T.ColorJitter(cj_bright, cj_contrast, cj_sat, cj_hue) - - # Transform for transformed jigsaw image - transform = [ - T.RandomHorizontalFlip(p=hf_prob), - T.RandomApply([color_jitter], p=cj_prob), - T.RandomGrayscale(p=random_gray_scale), - T.ToTensor(), - ] - - if normalize: - transform += [T.Normalize(mean=normalize["mean"], std=normalize["std"])] - - # Cropping and normalisation for untransformed image - self.no_augment = T.Compose( - [ - T.RandomResizedCrop(size=input_size, scale=(min_scale, 1.0)), - T.ToTensor(), - T.Normalize(mean=normalize["mean"], std=normalize["std"]), - ] - ) - self.jigsaw = Jigsaw( - n_grid=n_grid, - img_size=input_size_, - crop_size=int(input_size_ // n_grid), - transform=T.Compose(transform), - ) - - def forward(self, batch: List[tuple]): - """Overriding the BaseCollateFunction class's forward method because - for PIRL we need only one augmented batch, as opposed to both, which the - BaseCollateFunction creates. - """ - batch_size = len(batch) - - # list of transformed images - img_transforms = [ - self.jigsaw(batch[i][0]).unsqueeze_(0) for i in range(batch_size) - ] - img = [self.no_augment(batch[i][0]).unsqueeze_(0) for i in range(batch_size)] - # list of labels - labels = torch.LongTensor([item[1] for item in batch]) - # list of filenames - fnames = [item[2] for item in batch] - - # tuple of transforms - transforms = (torch.cat(img, 0), torch.cat(img_transforms, 0)) - - return transforms, labels, fnames - - -class MSNCollateFunction(MultiViewCollateFunction): - """Implements the transformations for MSN [0]. - - Generates a set of random and focal views for each input image. The generated output - is (views, target, filenames) where views is list with the following entries: - [random_views_0, random_views_1, ..., focal_views_0, focal_views_1, ...]. - - - [0]: Masked Siamese Networks, 2022: https://arxiv.org/abs/2204.07141 - - Attributes: - random_size: - Size of the random image views in pixels. - focal_size: - Size of the focal image views in pixels. - random_views: - Number of random views to generate. - focal_views: - Number of focal views to generate. - random_crop_scale: - Minimum and maximum size of the randomized crops for the relative to random_size. - focal_crop_scale: - Minimum and maximum size of the randomized crops relative to focal_size. - cj_prob: - Probability that color jittering is applied. - cj_strength: - Strength of the color jitter. - gaussian_blur: - Probability of Gaussian blur. - kernel_size: - Will be deprecated in favor of `sigmas` argument. If set, the old behavior applies and `sigmas` is ignored. - Used to calculate sigma of gaussian blur with kernel_size * input_size. - sigmas: - Tuple of min and max value from which the std of the gaussian kernel is sampled. - Is ignored if `kernel_size` is set. - random_gray_scale: - Probability of conversion to grayscale. - hf_prob: - Probability that horizontal flip is applied. - vf_prob: - Probability that vertical flip is applied. - normalize: - Dictionary with 'mean' and 'std' for torchvision.transforms.Normalize. - """ - - def __init__( - self, - random_size: int = 224, - focal_size: int = 96, - random_views: int = 2, - focal_views: int = 10, - random_crop_scale: Tuple[float, float] = (0.3, 1.0), - focal_crop_scale: Tuple[float, float] = (0.05, 0.3), - cj_prob: float = 0.8, - cj_strength: float = 1.0, - gaussian_blur: float = 0.5, - kernel_size: Optional[float] = None, - sigmas: Tuple[float, float] = (0.2, 2), - random_gray_scale: float = 0.2, - hf_prob: float = 0.5, - vf_prob: float = 0.0, - normalize: dict = imagenet_normalize, - ) -> None: - color_jitter = T.ColorJitter( - brightness=0.8 * cj_strength, - contrast=0.8 * cj_strength, - saturation=0.8 * cj_strength, - hue=0.2 * cj_strength, - ) - transform = T.Compose( - [ - T.RandomResizedCrop(size=random_size, scale=random_crop_scale), - T.RandomHorizontalFlip(p=hf_prob), - T.RandomVerticalFlip(p=vf_prob), - T.RandomApply([color_jitter], p=cj_prob), - T.RandomGrayscale(p=random_gray_scale), - GaussianBlur( - kernel_size=kernel_size, sigmas=sigmas, prob=gaussian_blur - ), - T.ToTensor(), - T.Normalize(mean=normalize["mean"], std=normalize["std"]), - ] - ) - focal_transform = T.Compose( - [ - T.RandomResizedCrop(size=focal_size, scale=focal_crop_scale), - T.RandomHorizontalFlip(p=hf_prob), - T.RandomVerticalFlip(p=vf_prob), - T.RandomApply([color_jitter], p=cj_prob), - T.RandomGrayscale(p=random_gray_scale), - GaussianBlur( - kernel_size=kernel_size, sigmas=sigmas, prob=gaussian_blur - ), - T.ToTensor(), - T.Normalize(mean=normalize["mean"], std=normalize["std"]), - ] - ) - transforms = [transform] * random_views - transforms += [focal_transform] * focal_views - super().__init__(transforms=transforms) - - -class SMoGCollateFunction(MultiViewCollateFunction): - """Implements the transformations for SMoG. - - Attributes: - crop_sizes: - Size of the input image in pixels for each crop category. - crop_counts: - Number of crops for each crop category. - crop_min_scales: - Min scales for each crop category. - crop_max_scales: - Max_scales for each crop category. - gaussian_blur_probs: - Probability of Gaussian blur for each crop category. - gaussian_blur_kernel_sizes: - Deprecated values in favour of sigmas. - gaussian_blur_sigmas: - Tuple of min and max value from which the std of the gaussian kernel is sampled. - solarize_probs: - Probability of solarization for each crop category. - hf_prob: - Probability that horizontal flip is applied. - cj_prob: - Probability that color jitter is applied. - cj_strength: - Strength of the color jitter. - random_gray_scale: - Probability of conversion to grayscale. - normalize: - Dictionary with 'mean' and 'std' for torchvision.transforms.Normalize. - - """ - - def __init__( - self, - crop_sizes: List[int] = [224, 96], - crop_counts: List[int] = [4, 4], - crop_min_scales: List[float] = [0.2, 0.05], - crop_max_scales: List[float] = [1.0, 0.2], - gaussian_blur_probs: List[float] = [0.5, 0.1], - gaussian_blur_kernel_sizes: Optional[List[float]] = [None, None], - gaussian_blur_sigmas: Tuple[float, float] = (0.2, 2), - solarize_probs: List[float] = [0.0, 0.2], - hf_prob: float = 0.5, - cj_prob: float = 1.0, - cj_strength: float = 0.5, - random_gray_scale: float = 0.2, - normalize: dict = imagenet_normalize, - ): - transforms = [] - for i in range(len(crop_sizes)): - random_resized_crop = T.RandomResizedCrop( - crop_sizes[i], scale=(crop_min_scales[i], crop_max_scales[i]) - ) - - color_jitter = T.ColorJitter( - 0.8 * cj_strength, - 0.8 * cj_strength, - 0.4 * cj_strength, - 0.2 * cj_strength, - ) - - transforms.extend( - [ - T.Compose( - [ - random_resized_crop, - T.RandomHorizontalFlip(p=hf_prob), - T.RandomApply([color_jitter], p=cj_prob), - T.RandomGrayscale(p=random_gray_scale), - GaussianBlur( - kernel_size=gaussian_blur_kernel_sizes[i], - prob=gaussian_blur_probs[i], - sigmas=gaussian_blur_sigmas, - ), # TODO - RandomSolarization(prob=solarize_probs[i]), - T.ToTensor(), - T.Normalize(mean=normalize["mean"], std=normalize["std"]), - ] - ) - ] - * crop_counts[i] - ) - - super().__init__(transforms) - - -class VICRegCollateFunction(BaseCollateFunction): - """Implementation of a collate function for images. - - This is an implementation of the BaseCollateFunction with a concrete - set of transforms. - - The set of transforms is inspired by the SimCLR paper as it has shown - to produce powerful embeddings. - - Attributes: - input_size: - Size of the input image in pixels. - cj_prob: - Probability that color jitter is applied. - cj_bright: - How much to jitter brightness. - cj_contrast: - How much to jitter contrast. - cj_sat: - How much to jitter saturation. - cj_hue: - How much to jitter hue. - min_scale: - Minimum size of the randomized crop relative to the input_size. - random_gray_scale: - Probability of conversion to grayscale. - solarize_prob: - Probability of solarization. - gaussian_blur: - Probability of Gaussian blur. - kernel_size: - Will be deprecated in favor of `sigmas` argument. If set, the old behavior applies and `sigmas` is ignored. - Used to calculate sigma of gaussian blur with kernel_size * input_size. - sigmas: - Tuple of min and max value from which the std of the gaussian kernel is sampled. - Is ignored if `kernel_size` is set. - vf_prob: - Probability that vertical flip is applied. - hf_prob: - Probability that horizontal flip is applied. - rr_prob: - Probability that random rotation is applied. - rr_degrees: - Range of degrees to select from for random rotation. If rr_degrees is None, - images are rotated by 90 degrees. If rr_degrees is a (min, max) tuple, - images are rotated by a random angle in [min, max]. If rr_degrees is a - single number, images are rotated by a random angle in - [-rr_degrees, +rr_degrees]. All rotations are counter-clockwise. - normalize: - Dictionary with 'mean' and 'std' for torchvision.transforms.Normalize. - - """ - - def __init__( - self, - input_size: int = 224, - cj_prob: float = 0.8, - cj_bright: float = 0.4, - cj_contrast: float = 0.4, - cj_sat: float = 0.2, - cj_hue: float = 0.1, - min_scale: float = 0.08, - random_gray_scale: float = 0.2, - solarize_prob: float = 0.1, - gaussian_blur: float = 0.5, - kernel_size: Optional[float] = None, - sigmas: Tuple[float, float] = (0.2, 2), - vf_prob: float = 0.0, - hf_prob: float = 0.5, - rr_prob: float = 0.0, - rr_degrees: Optional[Union[float, Tuple[float, float]]] = None, - normalize: dict = imagenet_normalize, - ): - if isinstance(input_size, tuple): - input_size_ = max(input_size) - else: - input_size_ = input_size - - color_jitter = T.ColorJitter(cj_bright, cj_contrast, cj_sat, cj_hue) - - transform = [ - T.RandomResizedCrop(size=input_size, scale=(min_scale, 1.0)), - T.RandomHorizontalFlip(p=hf_prob), - T.RandomVerticalFlip(p=vf_prob), - random_rotation_transform(rr_prob=rr_prob, rr_degrees=rr_degrees), - T.RandomApply([color_jitter], p=cj_prob), - T.RandomGrayscale(p=random_gray_scale), - RandomSolarization(prob=solarize_prob), - GaussianBlur(kernel_size=kernel_size, sigmas=sigmas, prob=gaussian_blur), - T.ToTensor(), - ] - - if normalize: - transform += [T.Normalize(mean=normalize["mean"], std=normalize["std"])] - - transform = T.Compose(transform) - - super(VICRegCollateFunction, self).__init__(transform) - - -class VICRegLCollateFunction(nn.Module): - """Transforms images for VICRegL. - - Attributes: - global_crop_size: - Size of the input image in pixels for the global crop category. - local_crop_size: - Size of the input image in pixels for the local crop category. - global_crop_scale: - Min and max scales for the global crop category. - local_crop_scale: - Min and max scales for the local crop category. - global_grid_size: - Grid size for the global crop category. - local_grid_size: - Grid size for the local crop category. - global_gaussian_blur_prob: - Probability of Gaussian blur for the global crop category. - local_gaussian_blur_prob: - Probability of Gaussian blur for the local crop category. - global_gaussian_blur_kernel_size: - Will be deprecated in favor of `global_gaussian_blur_sigmas` argument. If set, the old behavior applies and `global_gaussian_blur_sigmas` is ignored. - Used to calculate sigma of gaussian blur with global_gaussian_blur_kernel_size * input_size. Applied to global crop category. - local_gaussian_blur_kernel_size: - Will be deprecated in favor of `local_gaussian_blur_sigmas` argument. If set, the old behavior applies and `local_gaussian_blur_sigmas` is ignored. - Used to calculate sigma of gaussian blur with local_gaussian_blur_kernel_size * input_size. Applied to local crop category. - global_gaussian_blur_sigmas: - Tuple of min and max value from which the std of the gaussian kernel is sampled. - Is ignored if `global_gaussian_blur_kernel_size` is set. Applied to global crop category. - local_gaussian_blur_sigmas: - Tuple of min and max value from which the std of the gaussian kernel is sampled. - Is ignored if `local_gaussian_blur_kernel_size` is set. Applied to local crop category. - global_solarize_prob: - Probability of solarization for the global crop category. - local_solarize_prob: - Probability of solarization for the local crop category. - hf_prob: - Probability that horizontal flip is applied. - cj_prob: - Probability that color jitter is applied. - cj_strength: - Strength of the color jitter. - random_gray_scale: - Probability of conversion to grayscale. - normalize: - Dictionary with mean and standard deviation for normalization. - """ - - def __init__( - self, - global_crop_size: int = 224, - local_crop_size: int = 96, - global_crop_scale: Tuple[int] = (0.2, 1.0), - local_crop_scale: Tuple[int] = (0.05, 0.2), - global_grid_size: int = 7, - local_grid_size: int = 3, - global_gaussian_blur_prob: float = 0.5, - local_gaussian_blur_prob: float = 0.1, - global_gaussian_blur_kernel_size: Optional[float] = None, - local_gaussian_blur_kernel_size: Optional[float] = None, - global_gaussian_blur_sigmas: Tuple[float, float] = (0.2, 2), - local_gaussian_blur_sigmas: Tuple[float, float] = (0.2, 2), - global_solarize_prob: float = 0.0, - local_solarize_prob: float = 0.2, - hf_prob: float = 0.5, - cj_prob: float = 1.0, - cj_strength: float = 0.5, - random_gray_scale: float = 0.2, - normalize: dict = imagenet_normalize, - ): - _deprecation_warning_collate_functions() - super().__init__() - self.global_crop_and_flip = RandomResizedCropAndFlip( - crop_size=global_crop_size, - crop_min_scale=global_crop_scale[0], - crop_max_scale=global_crop_scale[1], - hf_prob=hf_prob, - grid_size=global_grid_size, - ) - self.local_crop_and_flip = RandomResizedCropAndFlip( - crop_size=local_crop_size, - crop_min_scale=local_crop_scale[0], - crop_max_scale=local_crop_scale[1], - hf_prob=hf_prob, - grid_size=local_grid_size, - ) - - color_jitter = T.ColorJitter( - 0.8 * cj_strength, - 0.8 * cj_strength, - 0.4 * cj_strength, - 0.2 * cj_strength, - ) - self.global_transform = T.Compose( - [ - T.RandomApply([color_jitter], p=cj_prob), - T.RandomGrayscale(p=random_gray_scale), - GaussianBlur( - kernel_size=global_gaussian_blur_kernel_size, - prob=global_gaussian_blur_prob, - sigmas=global_gaussian_blur_sigmas, - ), - RandomSolarization(prob=global_solarize_prob), - T.ToTensor(), - T.Normalize(mean=normalize["mean"], std=normalize["std"]), - ] - ) - - self.local_transform = T.Compose( - [ - T.RandomApply([color_jitter], p=cj_prob), - T.RandomGrayscale(p=random_gray_scale), - GaussianBlur( - kernel_size=local_gaussian_blur_kernel_size, - prob=local_gaussian_blur_prob, - sigmas=local_gaussian_blur_sigmas, - ), - RandomSolarization(prob=local_solarize_prob), - T.ToTensor(), - T.Normalize(mean=normalize["mean"], std=normalize["std"]), - ] - ) - - def forward( - self, batch: List[Tuple[Image.Image, int, str]] - ) -> Tuple[ - Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor], - torch.Tensor, - torch.Tensor, - ]: - """Applies transforms to images in the input batch. - - Args: - batch: - A list of tuples containing an image (as a PIL Image), - a label (int), and a filename (str). - - Returns: - A tuple of transformed images (as a 4-tuple of torch.Tensors containing view_global, view_local, grid_global, grid_local), - labels (as torch.Tensor), and filenames (as torch.Tensor). - - """ - views_global = [] - views_local = [] - grids_global = [] - grids_local = [] - labels = [] - fnames = [] - - for image, label, filename in batch: - view_global, grid_global = self.global_crop_and_flip.forward(image) - view_local, grid_local = self.local_crop_and_flip.forward(image) - views_global.append(self.global_transform(view_global)) - views_local.append(self.local_transform(view_local)) - grids_global.append(grid_global) - grids_local.append(grid_local) - labels.append(torch.LongTensor(label)) - fnames.append(filename) - - views_global = torch.stack(views_global) - views_local = torch.stack(views_local) - grids_global = torch.stack(grids_global) - grids_local = torch.stack(grids_local) - - return (views_global, views_local, grids_global, grids_local), labels, fnames - - -def _deprecation_warning_collate_functions() -> None: - warn( - "Collate functions are deprecated and will be removed in favor of transforms in v1.4.0.\n" - "See https://docs.lightly.ai/self-supervised-learning/examples/models.html for examples.", - category=DeprecationWarning, - ) - - -if TYPE_CHECKING: - # IJEPAMaskCollator moved to lightly.data.ijepa_collate. Bind the name for type - # checkers only; at runtime __getattr__ below serves it with a warning. - from lightly.data.ijepa_collate import IJEPAMaskCollator -else: - - def __getattr__(name: str) -> Any: - # Resolve the moved name lazily (PEP 562) so the old import path warns - # instead of failing, and raises AttributeError once it is removed. - if name == "IJEPAMaskCollator": - from lightly.data.ijepa_collate import IJEPAMaskCollator - from lightly.utils.deprecation import warn_deprecated - - warn_deprecated( - "Importing IJEPAMaskCollator from lightly.data.collate", - "lightly.data.IJEPAMaskCollator", - removed_in="1.7.0", - ) - return IJEPAMaskCollator - raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/lightly/utils/debug.py b/lightly/utils/debug.py index 60211ca73..5204d86ba 100644 --- a/lightly/utils/debug.py +++ b/lightly/utils/debug.py @@ -1,10 +1,10 @@ -from typing import List, Union +from typing import List import torch import torchvision from PIL import Image -from lightly.data.collate import BaseCollateFunction, MultiViewCollateFunction +from lightly.transforms.multi_view_transform import MultiViewTransform from lightly.transforms.torchvision_v2_compatibility import torchvision_transforms as T try: @@ -67,17 +67,26 @@ def apply_transform_without_normalize( type(T.ToTensor()), T.Normalize, ) + if isinstance(transform, skippable_transforms): + # Checked first because our ToTensor shim also wraps a T.Compose, which the + # branch below would otherwise unwrap and apply. + return image if isinstance(transform, T.Compose): for transform_ in transform.transforms: image = apply_transform_without_normalize(image, transform_) - elif not isinstance(transform, skippable_transforms): + elif isinstance(getattr(transform, "transform", None), T.Compose): + # View transforms such as SimCLRViewTransform wrap a T.Compose instead of + # subclassing it. torchvision v2's Compose has a bound method named + # "transform", so the isinstance check above has to come first. + image = apply_transform_without_normalize(image, transform.transform) + else: image = transform(image) return image def generate_grid_of_augmented_images( input_images: List[Image.Image], - collate_function: Union[BaseCollateFunction, MultiViewCollateFunction], + transform: MultiViewTransform, ) -> List[List[Image.Image]]: """Returns a grid of augmented images. Images in a column belong together. @@ -86,43 +95,32 @@ def generate_grid_of_augmented_images( Args: input_images: List of PIL images for which the augmentations should be plotted. - collate_function: - The collate function of the self-supervised learning algorithm. - Must be of type BaseCollateFunction or MultiViewCollateFunction. + transform: + The transform of the self-supervised learning algorithm. Must be of type + MultiViewTransform. Returns: A grid of augmented images. Images in a column belong together. """ - grid = [] - if isinstance(collate_function, BaseCollateFunction): - for _ in range(2): - grid.append( - [ - apply_transform_without_normalize(image, collate_function.transform) - for image in input_images - ] - ) - elif isinstance(collate_function, MultiViewCollateFunction): - for transform in collate_function.transforms: - grid.append( - [ - apply_transform_without_normalize(image, transform) - for image in input_images - ] - ) - else: + if not isinstance(transform, MultiViewTransform): raise ValueError( - "Collate function must be one of " - "(BaseCollateFunction, MultiViewCollateFunction) " - f"but is {type(collate_function)}." + f"Transform must be of type MultiViewTransform but is {type(transform)}." + ) + grid = [] + for view_transform in transform.transforms: + grid.append( + [ + apply_transform_without_normalize(image, view_transform) + for image in input_images + ] ) return grid def plot_augmented_images( input_images: List[Image.Image], - collate_function: Union[BaseCollateFunction, MultiViewCollateFunction], + transform: MultiViewTransform, ): """Plots original images and augmented images in a figure. @@ -131,15 +129,13 @@ def plot_augmented_images( Args: input_images: List of PIL images for which the augmentations should be plotted. - collate_function: - The collate function of the self-supervised learning algorithm. - Must be of type BaseCollateFunction or MultiViewCollateFunction. + transform: + The transform of the self-supervised learning algorithm. Must be of type + MultiViewTransform. Returns: A figure showing the original images in the left column and the augmented - images to their right. If the collate_function is an instance of the - BaseCollateFunction, two example augmentations are shown. For - MultiViewCollateFunctions all the generated views are shown. + images to their right. All views generated by the transform are shown. """ _check_matplotlib_available() @@ -147,7 +143,7 @@ def plot_augmented_images( if len(input_images) == 0: raise ValueError("There must be at least one input image.") - grid = generate_grid_of_augmented_images(input_images, collate_function) + grid = generate_grid_of_augmented_images(input_images, transform) grid.insert(0, input_images) nrows = len(input_images) ncols = len(grid) diff --git a/pyproject.toml b/pyproject.toml index ebbf26e7d..11283481d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -192,7 +192,6 @@ exclude = '''(?x)( lightly/cli/_cli_simclr.py | lightly/cli/_helpers.py | lightly/data/dataset.py | - lightly/data/collate.py | lightly/data/_image_loaders.py | lightly/data/_video.py | lightly/core.py | @@ -212,7 +211,6 @@ exclude = '''(?x)( tests/cli/test_cli_embed.py | tests/core/test_Core.py | tests/data/test_multi_view_collate.py | - tests/data/test_data_collate.py | tests/data/test_LightlySubset.py | tests/data/test_LightlyDataset.py | tests/embedding/test_callbacks.py | diff --git a/tests/data/test_data_collate.py b/tests/data/test_data_collate.py deleted file mode 100644 index 67e1351dc..000000000 --- a/tests/data/test_data_collate.py +++ /dev/null @@ -1,239 +0,0 @@ -import random - -import pytest -import torch -import torchvision - -from lightly.data import ( - BaseCollateFunction, - ImageCollateFunction, - MultiCropCollateFunction, - PIRLCollateFunction, - SimCLRCollateFunction, - SwaVCollateFunction, -) -from lightly.data.collate import ( - DINOCollateFunction, - MAECollateFunction, - MSNCollateFunction, - MultiViewCollateFunction, - VICRegCollateFunction, - VICRegLCollateFunction, -) -from lightly.transforms.torchvision_v2_compatibility import torchvision_transforms as T - - -class TestDataCollate: - def create_batch(self, batch_size=16, seed=0): - torch.manual_seed(0) - rnd_images = torchvision.datasets.FakeData(size=batch_size) - - fnames = [f"img_{i}.jpg" for i in range(batch_size)] - labels = [random.randint(0, 5) for i in range(batch_size)] - - batch = [] - - for i in range(batch_size): - batch.append((rnd_images[i][0], labels[i], fnames[i])) - - return batch - - def test_base_collate(self) -> None: - batch = self.create_batch() - transform = T.ToTensor() - collate = BaseCollateFunction(transform) - samples, labels, fnames = collate(batch) - samples0, samples1 = samples - - assert collate is not None - assert len(samples0) == len(samples1) - assert len(samples1) == len(labels) - - def test_image_collate(self) -> None: - batch = self.create_batch() - img_collate = ImageCollateFunction() - samples, labels, fnames = img_collate(batch) - samples0, samples1 = samples - - assert img_collate is not None - assert len(samples0) == len(samples1) - assert len(samples1) == len(labels) - - def test_image_collate_tuple_input_size(self) -> None: - batch = self.create_batch() - img_collate = ImageCollateFunction( - input_size=(32, 32), - ) - samples, labels, fnames = img_collate(batch) - samples0, samples1 = samples - - assert img_collate is not None - assert len(samples0) == len(samples1) - assert len(samples1) == len(labels) - - def test_image_collate_random_rotate(self) -> None: - batch = self.create_batch() - img_collate = ImageCollateFunction(rr_prob=1.0, rr_degrees=45.0) - samples, labels, fnames = img_collate(batch) - samples0, samples1 = samples - - assert img_collate is not None - assert len(samples0) == len(samples1) - assert len(samples1) == len(labels) - - def test_image_collate_random_rotate__tuple_degrees(self) -> None: - batch = self.create_batch() - img_collate = ImageCollateFunction(rr_prob=1.0, rr_degrees=(-15.0, 45.0)) - samples, labels, fnames = img_collate(batch) - samples0, samples1 = samples - - assert img_collate is not None - assert len(samples0) == len(samples1) - assert len(samples1) == len(labels) - - def test_simclr_collate_tuple_input_size(self) -> None: - batch = self.create_batch() - img_collate = SimCLRCollateFunction( - input_size=(32, 32), - ) - samples, labels, fnames = img_collate(batch) - samples0, samples1 = samples - - assert img_collate is not None - assert len(samples0) == len(samples1) - assert len(samples1) == len(labels) - - @pytest.mark.parametrize("low", range(6)) - @pytest.mark.parametrize("high", range(2, 4)) - def test_multi_crop_collate(self, high: int, low: int) -> None: - batch = self.create_batch() - multi_crop_collate = MultiCropCollateFunction( - crop_sizes=[32, 16], - crop_counts=[high, low], - crop_min_scales=[0.14, 0.04], - crop_max_scales=[1.0, 0.14], - transforms=T.ToTensor(), - ) - samples, labels, fnames = multi_crop_collate(batch) - assert multi_crop_collate is not None - assert len(samples) == low + high - for i, crop in enumerate(samples): - if i < high: - assert crop.shape[-1] == 32 - assert crop.shape[-2] == 32 - else: - assert crop.shape[-1] == 16 - assert crop.shape[-2] == 16 - assert len(crop) == len(labels) - - def test_swav_collate_init(self) -> None: - swav_collate = SwaVCollateFunction() - - def test_swav_collate_init_fail(self) -> None: - with pytest.raises(ValueError): - SwaVCollateFunction( - crop_sizes=[1], - crop_counts=[2, 3], - ) - - def test_multi_view_collate(self) -> None: - to_tensor = T.ToTensor() - hflip = T.Compose( - [ - T.RandomHorizontalFlip(p=1), - to_tensor, - ] - ) - vflip = T.Compose( - [ - T.RandomVerticalFlip(p=1), - to_tensor, - ] - ) - trans = [to_tensor, hflip, vflip] - - collate_fn = MultiViewCollateFunction(trans) - batch = self.create_batch() - imgs = batch[0] - views, labels, fnames = collate_fn(batch) - - assert len(labels) == len(batch) - assert len(fnames) == len(batch) - assert torch.equal(views[0][0], to_tensor(imgs[0])) - assert torch.equal(views[1][0], hflip(imgs[0])) - assert torch.equal(views[2][0], vflip(imgs[0])) - - def test_dino_collate_init(self) -> None: - DINOCollateFunction() - - def test_dino_collate_forward(self) -> None: - batch = self.create_batch() - collate_fn = DINOCollateFunction() - views, labels, fnames = collate_fn(batch) - - def test_mae_collate_init(self) -> None: - MAECollateFunction() - - def test_mae_collate_forward(self) -> None: - batch = self.create_batch() - collate_fn = MAECollateFunction() - views, labels, fnames = collate_fn(batch) - - def test_pirl_collate_init(self) -> None: - PIRLCollateFunction() - - def test_pirl_collate_forward_tuple_input_size(self) -> None: - batch = self.create_batch() - img_collate = PIRLCollateFunction( - input_size=(32, 32), - ) - samples, labels, fnames = img_collate(batch) - samples0, samples1 = samples - - assert img_collate is not None - assert len(samples0) == len(samples1) - assert len(samples1) == len(labels) - - def test_pirl_collate_forward_n_grid(self) -> None: - batch = self.create_batch() - img_collate = PIRLCollateFunction(input_size=32, n_grid=3) - samples, labels, fnames = img_collate(batch) - samples0, samples1 = samples - - assert img_collate is not None - assert len(samples0) == len(samples1) - assert len(samples1) == len(labels) - assert samples1.shape == (16, 9, 3, 10, 10) - - def test_msn_collate_init(self) -> None: - MSNCollateFunction() - - def test_msn_collate_forward(self) -> None: - batch = self.create_batch() - img_collate = MSNCollateFunction( - random_size=24, focal_size=12, random_views=2, focal_views=10 - ) - views, labels, fnames = img_collate(batch) - assert len(views) == 2 + 10 - assert len(labels) == len(batch) - assert len(fnames) == len(batch) - for view in views[:2]: - assert view.shape == (16, 3, 24, 24) - for view in views[2:]: - assert view.shape == (16, 3, 12, 12) - - def test_vicreg_collate_init(self) -> None: - VICRegCollateFunction() - - def test_vicreg_collate_forward(self) -> None: - batch = self.create_batch() - collate_fn = VICRegCollateFunction() - views, labels, fnames = collate_fn(batch) - - def test_vicregl_collate_init(self) -> None: - VICRegLCollateFunction() - - def test_vicregl_collate_forward(self) -> None: - batch = self.create_batch() - collate_fn = VICRegLCollateFunction() - views, labels, fnames = collate_fn(batch) diff --git a/tests/data/test_ijepa_collate.py b/tests/data/test_ijepa_collate.py index f5beda9f8..ca147761a 100644 --- a/tests/data/test_ijepa_collate.py +++ b/tests/data/test_ijepa_collate.py @@ -58,15 +58,6 @@ def test_call__allow_overlap() -> None: assert len(masks_pred) == 2 -def test_collate_reexport__warns() -> None: - with pytest.warns(FutureWarning, match="lightly.data.collate"): - from lightly.data.collate import IJEPAMaskCollator as reexported - - assert reexported is IJEPAMaskCollator - - -def test_collate_reexport__unknown_attribute() -> None: - import lightly.data.collate - - with pytest.raises(AttributeError): - getattr(lightly.data.collate, "DoesNotExist") +def test_collate_module__removed() -> None: + with pytest.raises(ImportError): + import lightly.data.collate # noqa: F401 diff --git a/tests/utils/test_debug.py b/tests/utils/test_debug.py index b3f6e2355..e0db7d827 100644 --- a/tests/utils/test_debug.py +++ b/tests/utils/test_debug.py @@ -5,7 +5,12 @@ import torch from PIL import Image -from lightly.data import collate +from lightly.transforms import ( + BYOLTransform, + DINOTransform, + SimCLRTransform, + SwaVTransform, +) from lightly.utils import debug try: @@ -42,33 +47,52 @@ def test_std_of_l2_normalized_raises(self): debug.std_of_l2_normalized(z) @pytest.mark.skipif(not MATPLOTLIB_AVAILABLE, reason="Matplotlib not installed") - def test_plot_augmented_images_image_collate_function(self): - # simclr collate function is a subclass of the image collate function - collate_function = collate.SimCLRCollateFunction() + def test_plot_augmented_images(self): + transform = SimCLRTransform(input_size=32) for n_images in range(2, 10): images = [self._generate_random_image(100, 100, 3) for _ in range(n_images)] - fig = debug.plot_augmented_images(images, collate_function) + fig = debug.plot_augmented_images(images, transform) assert fig is not None @pytest.mark.skipif(not MATPLOTLIB_AVAILABLE, reason="Matplotlib not installed") - def test_plot_augmented_images_multi_view_collate_function(self): - # dion collate function is a subclass of the multi view collate function - collate_function = collate.DINOCollateFunction() + def test_plot_augmented_images_many_views(self): + transform = DINOTransform(global_crop_size=32, local_crop_size=16) for n_images in range(1, 10): images = [self._generate_random_image(100, 100, 3) for _ in range(n_images)] - fig = debug.plot_augmented_images(images, collate_function) + fig = debug.plot_augmented_images(images, transform) assert fig is not None @pytest.mark.skipif(not MATPLOTLIB_AVAILABLE, reason="Matplotlib not installed") def test_plot_augmented_images_no_images(self): - collate_function = collate.SimCLRCollateFunction() with pytest.raises(ValueError): - debug.plot_augmented_images([], collate_function) + debug.plot_augmented_images([], SimCLRTransform(input_size=32)) @pytest.mark.skipif(not MATPLOTLIB_AVAILABLE, reason="Matplotlib not installed") - def test_plot_augmented_images_invalid_collate_function(self): + def test_plot_augmented_images_invalid_transform(self): images = [self._generate_random_image(100, 100, 3)] with pytest.raises(ValueError): debug.plot_augmented_images(images, None) + + @pytest.mark.parametrize( + "transform", + [ + SimCLRTransform(input_size=32), + DINOTransform(global_crop_size=32, local_crop_size=16), + SwaVTransform(crop_sizes=(32, 16)), + BYOLTransform(), + ], + ) + def test_generate_grid_of_augmented_images__returns_pil_images(self, transform): + # ToTensor and Normalize must be skipped, otherwise the grid holds tensors + # and plotting fails downstream. + images = [self._generate_random_image(100, 100, 3) for _ in range(2)] + + grid = debug.generate_grid_of_augmented_images(images, transform) + + assert len(grid) == len(transform.transforms) + for row in grid: + assert len(row) == len(images) + for image in row: + assert isinstance(image, Image.Image) From bd37c53c6704ae48553c024a469cc962b0797f27 Mon Sep 17 00:00:00 2001 From: gabrielfruet Date: Tue, 11 Aug 2026 18:07:28 -0300 Subject: [PATCH 2/9] Remove deprecated high-level model wrappers BarlowTwins, BYOL, MoCo, NNCLR, SimCLR and SimSiam have warned since they were marked for removal in 1.3.0. Build models from the heads in lightly.models.modules instead, as the examples do. _momentum.py goes with them; BYOL and MoCo were its only users. Co-Authored-By: Claude Opus 5 --- lightly/models/__init__.py | 10 - lightly/models/_momentum.py | 99 -------- lightly/models/barlowtwins.py | 122 ---------- lightly/models/byol.py | 167 ------------- lightly/models/moco.py | 136 ----------- lightly/models/nnclr.py | 227 ------------------ lightly/models/simclr.py | 108 --------- lightly/models/simsiam.py | 134 ----------- .../utils/benchmarking/benchmark_module.py | 18 +- pyproject.toml | 29 +-- tests/models/test_ModelsBYOL.py | 112 --------- tests/models/test_ModelsMoCo.py | 105 -------- tests/models/test_ModelsNNCLR.py | 126 ---------- tests/models/test_ModelsSimCLR.py | 105 -------- tests/models/test_ModelsSimSiam.py | 108 --------- 15 files changed, 12 insertions(+), 1594 deletions(-) delete mode 100644 lightly/models/_momentum.py delete mode 100644 lightly/models/barlowtwins.py delete mode 100644 lightly/models/byol.py delete mode 100644 lightly/models/moco.py delete mode 100644 lightly/models/nnclr.py delete mode 100644 lightly/models/simclr.py delete mode 100644 lightly/models/simsiam.py delete mode 100644 tests/models/test_ModelsBYOL.py delete mode 100644 tests/models/test_ModelsMoCo.py delete mode 100644 tests/models/test_ModelsNNCLR.py delete mode 100644 tests/models/test_ModelsSimCLR.py delete mode 100644 tests/models/test_ModelsSimSiam.py diff --git a/lightly/models/__init__.py b/lightly/models/__init__.py index c2f4e39c6..2a6de9841 100644 --- a/lightly/models/__init__.py +++ b/lightly/models/__init__.py @@ -1,9 +1,5 @@ """The lightly.models package provides model implementations. -Note that the high-level building blocks will be deprecated with -lightly version 1.3.0. Instead, use low-level building blocks to build the -models yourself. - Example implementations for all models can be found here: `Model Examples `_ @@ -19,11 +15,5 @@ # All Rights Reserved from lightly.models import utils -from lightly.models.barlowtwins import BarlowTwins -from lightly.models.byol import BYOL -from lightly.models.moco import MoCo -from lightly.models.nnclr import NNCLR from lightly.models.resnet import ResNetGenerator -from lightly.models.simclr import SimCLR -from lightly.models.simsiam import SimSiam from lightly.models.zoo import ZOO, checkpoints diff --git a/lightly/models/_momentum.py b/lightly/models/_momentum.py deleted file mode 100644 index b3fa9852b..000000000 --- a/lightly/models/_momentum.py +++ /dev/null @@ -1,99 +0,0 @@ -"""Momentum Encoder""" - -# Copyright (c) 2020. Lightly AG and its affiliates. -# All Rights Reserved - -import copy -from typing import Iterable, Tuple - -import torch -import torch.nn as nn -from torch import Tensor -from torch.nn.parameter import Parameter - - -def _deactivate_requires_grad(params: Iterable[Parameter]) -> None: - """Deactivates the requires_grad flag for all parameters.""" - for param in params: - param.requires_grad = False - - -def _do_momentum_update( - prev_params: Iterable[Parameter], params: Iterable[Parameter], m: float -) -> None: - """Updates the weights of the previous parameters.""" - for prev_param, param in zip(prev_params, params): - prev_param.data = prev_param.data * m + param.data * (1.0 - m) - - -class _MomentumEncoderMixin: - """Mixin to provide momentum encoder functionalities. - - Provides the following functionalities: - - Momentum encoder initialization. - - Momentum updates. - - Batch shuffling and unshuffling. - - To make use of the mixin, simply inherit from it: - - >>> class MyMoCo(nn.Module, _MomentumEncoderMixin): - >>> - >>> def __init__(self, backbone): - >>> super(MyMoCo, self).__init__() - >>> - >>> self.backbone = backbone - >>> self.projection_head = get_projection_head() - >>> - >>> self._init_momentum_encoder() # initialize momentum_backbone and momentum_projection_head - >>> - >>> def forward(self, x: Tensor): - >>> self._momentum_update(0.999) # do the momentum update - >>> - >>> y = self.momentum_backbone(x) # use momentum backbone - >>> y = self.momentum_projection_head(y) - - """ - - m: float - backbone: nn.Module - projection_head: nn.Module - momentum_backbone: nn.Module - momentum_projection_head: nn.Module - - def _init_momentum_encoder(self) -> None: - """Initializes momentum backbone and a momentum projection head.""" - assert self.backbone is not None - assert self.projection_head is not None - - self.momentum_backbone = copy.deepcopy(self.backbone) - self.momentum_projection_head = copy.deepcopy(self.projection_head) - - _deactivate_requires_grad(self.momentum_backbone.parameters()) - _deactivate_requires_grad(self.momentum_projection_head.parameters()) - - @torch.no_grad() - def _momentum_update(self, m: float = 0.999) -> None: - """Performs the momentum update for the backbone and projection head.""" - _do_momentum_update( - self.momentum_backbone.parameters(), - self.backbone.parameters(), - m=m, - ) - _do_momentum_update( - self.momentum_projection_head.parameters(), - self.projection_head.parameters(), - m=m, - ) - - @torch.no_grad() - def _batch_shuffle(self, batch: Tensor) -> Tuple[Tensor, Tensor]: - """Returns the shuffled batch and the indices to undo.""" - batch_size = batch.shape[0] - shuffle = torch.randperm(batch_size, device=batch.device) - return batch[shuffle], shuffle - - @torch.no_grad() - def _batch_unshuffle(self, batch: Tensor, shuffle: Tensor) -> Tensor: - """Returns the unshuffled batch.""" - unshuffle = torch.argsort(shuffle) - return batch[unshuffle] diff --git a/lightly/models/barlowtwins.py b/lightly/models/barlowtwins.py deleted file mode 100644 index 6021cea8a..000000000 --- a/lightly/models/barlowtwins.py +++ /dev/null @@ -1,122 +0,0 @@ -"""Barlow Twins resnet-based Model [0] -[0] Zbontar,J. et.al. 2021. Barlow Twins... https://arxiv.org/abs/2103.03230 -""" - -# Copyright (c) 2020. Lightly AG and its affiliates. -# All Rights Reserved - -import warnings - -import torch -import torch.nn as nn - -from lightly.models.modules import BarlowTwinsProjectionHead - - -class BarlowTwins(nn.Module): - """Implementation of BarlowTwins[0] network. - - Recommended loss: :py:class:`lightly.loss.barlow_twins_loss.BarlowTwinsLoss` - - Default params are the ones explained in the original paper [0]. - [0] Zbontar,J. et.al. 2021. Barlow Twins... https://arxiv.org/abs/2103.03230 - - Attributes: - backbone: - Backbone model to extract features from images. - ResNet-50 in original paper [0]. - num_ftrs: - Dimension of the embedding (before the projection head). - proj_hidden_dim: - Dimension of the hidden layer of the projection head. This should - be the same size as `num_ftrs`. - out_dim: - Dimension of the output (after the projection head). - - """ - - def __init__( - self, - backbone: nn.Module, - num_ftrs: int = 2048, - proj_hidden_dim: int = 8192, - out_dim: int = 8192, - ): - super(BarlowTwins, self).__init__() - - self.backbone = backbone - self.num_ftrs = num_ftrs - self.proj_hidden_dim = proj_hidden_dim - self.out_dim = out_dim - - self.projection_mlp = BarlowTwinsProjectionHead( - num_ftrs, proj_hidden_dim, out_dim - ) - - warnings.warn( - Warning( - "The high-level building block BarlowTwins will be deprecated in version 1.3.0. " - + "Use low-level building blocks instead. " - + "See https://docs.lightly.ai/self-supervised-learning/lightly.models.html for more information" - ), - DeprecationWarning, - ) - - def forward( - self, x0: torch.Tensor, x1: torch.Tensor = None, return_features: bool = False - ): - """Forward pass through BarlowTwins. - - Extracts features with the backbone and applies the projection - head to the output space. If both x0 and x1 are not None, both will be - passed through the backbone and projection. If x1 is None, only x0 will - be forwarded. - Barlow Twins only implement a projection head unlike SimSiam. - - Args: - x0: - Tensor of shape bsz x channels x W x H. - x1: - Tensor of shape bsz x channels x W x H. - return_features: - Whether or not to return the intermediate features backbone(x). - - Returns: - The output projection of x0 and (if x1 is not None) - the output projection of x1. If return_features is - True, the output for each x is a tuple (out, f) where f are the - features before the projection head. - - Examples: - >>> # single input, single output - >>> out = model(x) - >>> - >>> # single input with return_features=True - >>> out, f = model(x, return_features=True) - >>> - >>> # two inputs, two outputs - >>> out0, out1 = model(x0, x1) - >>> - >>> # two inputs, two outputs with return_features=True - >>> (out0, f0), (out1, f1) = model(x0, x1, return_features=True) - """ - # forward pass first input - f0 = self.backbone(x0).flatten(start_dim=1) - out0 = self.projection_mlp(f0) - - # append features if requested - if return_features: - out0 = (out0, f0) - - if x1 is None: - return out0 - - # forward pass second input - f1 = self.backbone(x1).flatten(start_dim=1) - out1 = self.projection_mlp(f1) - - # append features if requested - if return_features: - out1 = (out1, f1) - - return out0, out1 diff --git a/lightly/models/byol.py b/lightly/models/byol.py deleted file mode 100644 index 20e92c20e..000000000 --- a/lightly/models/byol.py +++ /dev/null @@ -1,167 +0,0 @@ -"""BYOL Model""" - -# Copyright (c) 2021. Lightly AG and its affiliates. -# All Rights Reserved - -import warnings - -import torch -import torch.nn as nn - -from lightly.models._momentum import _MomentumEncoderMixin -from lightly.models.modules import BYOLProjectionHead - - -def _get_byol_mlp(num_ftrs: int, hidden_dim: int, out_dim: int): - """Returns a 2-layer MLP with batch norm on the hidden layer. - - Reference (12.03.2021) - https://arxiv.org/abs/2006.07733 - - """ - modules = [ - nn.Linear(num_ftrs, hidden_dim), - nn.BatchNorm1d(hidden_dim), - nn.ReLU(), - nn.Linear(hidden_dim, out_dim), - ] - return nn.Sequential(*modules) - - -class BYOL(nn.Module, _MomentumEncoderMixin): - """Implementation of the BYOL architecture. - - Attributes: - backbone: - Backbone model to extract features from images. - num_ftrs: - Dimension of the embedding (before the projection mlp). - hidden_dim: - Dimension of the hidden layer in the projection and prediction mlp. - out_dim: - Dimension of the output (after the projection/prediction mlp). - m: - Momentum for the momentum update of encoder. - """ - - def __init__( - self, - backbone: nn.Module, - num_ftrs: int = 2048, - hidden_dim: int = 4096, - out_dim: int = 256, - m: float = 0.9, - ): - super(BYOL, self).__init__() - - self.backbone = backbone - # the architecture of the projection and prediction head is the same - self.projection_head = BYOLProjectionHead(num_ftrs, hidden_dim, out_dim) - self.prediction_head = BYOLProjectionHead(out_dim, hidden_dim, out_dim) - self.momentum_backbone = None - self.momentum_projection_head = None - - self._init_momentum_encoder() - self.m = m - - warnings.warn( - Warning( - "The high-level building block BYOL will be deprecated in version 1.3.0. " - + "Use low-level building blocks instead. " - + "See https://docs.lightly.ai/self-supervised-learning/lightly.models.html for more information" - ), - DeprecationWarning, - ) - - def _forward(self, x0: torch.Tensor, x1: torch.Tensor = None): - """Forward pass through the encoder and the momentum encoder. - - Performs the momentum update, extracts features with the backbone and - applies the projection (and prediction) head to the output space. If - x1 is None, only x0 will be processed otherwise, x0 is processed with - the encoder and x1 with the momentum encoder. - - Args: - x0: - Tensor of shape bsz x channels x W x H. - x1: - Tensor of shape bsz x channels x W x H. - - Returns: - The output projection of x0 and (if x1 is not None) the output - projection of x1. - - Examples: - >>> # single input, single output - >>> out = model._forward(x) - >>> - >>> # two inputs, two outputs - >>> out0, out1 = model._forward(x0, x1) - - """ - - self._momentum_update(self.m) - - # forward pass of first input x0 - f0 = self.backbone(x0).flatten(start_dim=1) - z0 = self.projection_head(f0) - out0 = self.prediction_head(z0) - - if x1 is None: - return out0 - - # forward pass of second input x1 - with torch.no_grad(): - f1 = self.momentum_backbone(x1).flatten(start_dim=1) - out1 = self.momentum_projection_head(f1) - - return out0, out1 - - def forward( - self, x0: torch.Tensor, x1: torch.Tensor, return_features: bool = False - ): - """Symmetrizes the forward pass (see _forward). - - Performs two forward passes, once where x0 is passed through the encoder - and x1 through the momentum encoder and once the other way around. - - Note that this model currently requires two inputs for the forward pass - (x0 and x1) which correspond to the two augmentations. - Furthermore, `the return_features` argument does not work yet. - - Args: - x0: - Tensor of shape bsz x channels x W x H. - x1: - Tensor of shape bsz x channels x W x H. - - Returns: - A tuple out0, out1, where out0 and out1 are tuples containing the - predictions and projections of x0 and x1: out0 = (z0, p0) and - out1 = (z1, p1). - - Examples: - >>> # initialize the model and the loss function - >>> model = BYOL() - >>> criterion = SymNegCosineSimilarityLoss() - >>> - >>> # forward pass for two batches of transformed images x1 and x2 - >>> out0, out1 = model(x0, x1) - >>> loss = criterion(out0, out1) - - """ - - if x0 is None: - raise ValueError("x0 must not be None!") - if x1 is None: - raise ValueError("x1 must not be None!") - - if not all([s0 == s1 for s0, s1 in zip(x0.shape, x1.shape)]): - raise ValueError( - f"x0 and x1 must have same shape but got shapes {x0.shape} and {x1.shape}!" - ) - - p0, z1 = self._forward(x0, x1) - p1, z0 = self._forward(x1, x0) - - return (z0, p0), (z1, p1) diff --git a/lightly/models/moco.py b/lightly/models/moco.py deleted file mode 100644 index ecabb9602..000000000 --- a/lightly/models/moco.py +++ /dev/null @@ -1,136 +0,0 @@ -"""MoCo Model""" - -# Copyright (c) 2020. Lightly AG and its affiliates. -# All Rights Reserved - -import warnings - -import torch -import torch.nn as nn - -from lightly.models._momentum import _MomentumEncoderMixin -from lightly.models.modules import MoCoProjectionHead - - -class MoCo(nn.Module, _MomentumEncoderMixin): - """Implementation of the MoCo (Momentum Contrast)[0] architecture. - - Recommended loss: :py:class:`lightly.loss.ntx_ent_loss.NTXentLoss` with - a memory bank. - - [0] MoCo, 2020, https://arxiv.org/abs/1911.05722 - - Attributes: - backbone: - Backbone model to extract features from images. - num_ftrs: - Dimension of the embedding (before the projection head). - out_dim: - Dimension of the output (after the projection head). - m: - Momentum for momentum update of the key-encoder. - - """ - - def __init__( - self, - backbone: nn.Module, - num_ftrs: int = 32, - out_dim: int = 128, - m: float = 0.999, - batch_shuffle: bool = False, - ): - super(MoCo, self).__init__() - - self.backbone = backbone - self.projection_head = MoCoProjectionHead(num_ftrs, num_ftrs, out_dim) - self.momentum_features = None - self.momentum_projection_head = None - - self.m = m - self.batch_shuffle = batch_shuffle - - # initialize momentum features and momentum projection head - self._init_momentum_encoder() - - warnings.warn( - Warning( - "The high-level building block MoCo will be deprecated in version 1.3.0. " - + "Use low-level building blocks instead. " - + "See https://docs.lightly.ai/self-supervised-learning/lightly.models.html for more information" - ), - DeprecationWarning, - ) - - def forward( - self, x0: torch.Tensor, x1: torch.Tensor = None, return_features: bool = False - ): - """Embeds and projects the input image. - - Performs the momentum update, extracts features with the backbone and - applies the projection head to the output space. If both x0 and x1 are - not None, both will be passed through the backbone and projection head. - If x1 is None, only x0 will be forwarded. - - Args: - x0: - Tensor of shape bsz x channels x W x H. - x1: - Tensor of shape bsz x channels x W x H. - return_features: - Whether or not to return the intermediate features backbone(x). - - Returns: - The output projection of x0 and (if x1 is not None) the output - projection of x1. If return_features is True, the output for each x - is a tuple (out, f) where f are the features before the projection - head. - - Examples: - >>> # single input, single output - >>> out = model(x) - >>> - >>> # single input with return_features=True - >>> out, f = model(x, return_features=True) - >>> - >>> # two inputs, two outputs - >>> out0, out1 = model(x0, x1) - >>> - >>> # two inputs, two outputs with return_features=True - >>> (out0, f0), (out1, f1) = model(x0, x1, return_features=True) - - """ - self._momentum_update(self.m) - - # forward pass of first input x0 - f0 = self.backbone(x0).flatten(start_dim=1) - out0 = self.projection_head(f0) - - # append features if requested - if return_features: - out0 = (out0, f0) - - # return out0 if x1 is None - if x1 is None: - return out0 - - # forward pass of second input x1 - with torch.no_grad(): - # shuffle for batchnorm - if self.batch_shuffle: - x1, shuffle = self._batch_shuffle(x1) - - # run x1 through momentum encoder - f1 = self.momentum_backbone(x1).flatten(start_dim=1) - out1 = self.momentum_projection_head(f1).detach() - - # unshuffle for batchnorm - if self.batch_shuffle: - f1 = self._batch_unshuffle(f1, shuffle) - out1 = self._batch_unshuffle(out1, shuffle) - - # append features if requested - if return_features: - out1 = (out1, f1) - - return out0, out1 diff --git a/lightly/models/nnclr.py b/lightly/models/nnclr.py deleted file mode 100644 index 82e649b75..000000000 --- a/lightly/models/nnclr.py +++ /dev/null @@ -1,227 +0,0 @@ -"""NNCLR Model""" - -# Copyright (c) 2021. Lightly AG and its affiliates. -# All Rights Reserved - -import warnings - -import torch -import torch.nn as nn - -from lightly.models.modules import NNCLRPredictionHead, NNCLRProjectionHead - - -def _prediction_mlp(in_dims: int, h_dims: int, out_dims: int) -> nn.Sequential: - """Prediction MLP. The original paper's implementation has 2 layers, with - BN applied to its hidden fc layers but no BN or ReLU on the output fc layer. - - Note that the hidden dimensions should be smaller than the input/output - dimensions (bottleneck structure). The default implementation using a - ResNet50 backbone has an input dimension of 2048, hidden dimension of 512, - and output dimension of 2048 - - Args: - in_dims: - Input dimension of the first linear layer. - h_dims: - Hidden dimension of all the fully connected layers (should be a - bottleneck!) - out_dims: - Output Dimension of the final linear layer. - - Returns: - nn.Sequential: - The projection head. - """ - l1 = nn.Sequential( - nn.Linear(in_dims, h_dims), nn.BatchNorm1d(h_dims), nn.ReLU(inplace=True) - ) - - l2 = nn.Linear(h_dims, out_dims) - - prediction = nn.Sequential(l1, l2) - return prediction - - -def _projection_mlp( - num_ftrs: int, h_dims: int, out_dim: int, num_layers: int = 3 -) -> nn.Sequential: - """Projection MLP. The original paper's implementation has 3 layers, with - BN applied to its hidden fc layers but no ReLU on the output fc layer. - The CIFAR-10 study used a MLP with only two layers. - - Args: - in_dims: - Input dimension of the first linear layer. - h_dims: - Hidden dimension of all the fully connected layers. - out_dims: - Output Dimension of the final linear layer. - num_layers: - Controls the number of layers; must be 2 or 3. Defaults to 3. - - Returns: - nn.Sequential: - The projection head. - """ - l1 = nn.Sequential( - nn.Linear(num_ftrs, h_dims), nn.BatchNorm1d(h_dims), nn.ReLU(inplace=True) - ) - - l2 = nn.Sequential( - nn.Linear(h_dims, h_dims), nn.BatchNorm1d(h_dims), nn.ReLU(inplace=True) - ) - - l3 = nn.Sequential(nn.Linear(h_dims, out_dim), nn.BatchNorm1d(out_dim)) - - if num_layers == 3: - projection = nn.Sequential(l1, l2, l3) - elif num_layers == 2: - projection = nn.Sequential(l1, l3) - else: - raise NotImplementedError("Only MLPs with 2 and 3 layers are implemented.") - - return projection - - -class NNCLR(nn.Module): - """Implementation of the NNCLR[0] architecture - - Recommended loss: :py:class:`lightly.loss.ntx_ent_loss.NTXentLoss` - Recommended module: :py:class:`lightly.models.modules.nn_memory_bank.NNmemoryBankModule` - - [0] NNCLR, 2021, https://arxiv.org/abs/2104.14548 - - Attributes: - backbone: - Backbone model to extract features from images. - num_ftrs: - Dimension of the embedding (before the projection head). - proj_hidden_dim: - Dimension of the hidden layer of the projection head. - pred_hidden_dim: - Dimension of the hidden layer of the predicion head. - out_dim: - Dimension of the output (after the projection head). - num_mlp_layers: - Number of linear layers for MLP. - - Examples: - >>> model = NNCLR(backbone) - >>> criterion = NTXentLoss(temperature=0.1) - >>> - >>> nn_replacer = NNmemoryBankModule(size=2**16) - >>> - >>> # forward pass - >>> (z0, p0), (z1, p1) = model(x0, x1) - >>> z0 = nn_replacer(z0.detach(), update=False) - >>> z1 = nn_replacer(z1.detach(), update=True) - >>> - >>> loss = 0.5 * (criterion(z0, p1) + criterion(z1, p0)) - - """ - - def __init__( - self, - backbone: nn.Module, - num_ftrs: int = 512, - proj_hidden_dim: int = 2048, - pred_hidden_dim: int = 4096, - out_dim: int = 256, - ): - super(NNCLR, self).__init__() - - self.backbone = backbone - self.num_ftrs = num_ftrs - self.proj_hidden_dim = proj_hidden_dim - self.pred_hidden_dim = pred_hidden_dim - self.out_dim = out_dim - - self.projection_mlp = NNCLRProjectionHead( - num_ftrs, - proj_hidden_dim, - out_dim, - ) - - self.prediction_mlp = NNCLRPredictionHead( - out_dim, - pred_hidden_dim, - out_dim, - ) - - warnings.warn( - Warning( - "The high-level building block NNCLR will be deprecated in version 1.3.0. " - + "Use low-level building blocks instead. " - + "See https://docs.lightly.ai/self-supervised-learning/lightly.models.html for more information" - ), - DeprecationWarning, - ) - - def forward( - self, x0: torch.Tensor, x1: torch.Tensor = None, return_features: bool = False - ): - """Embeds and projects the input images. - - Extracts features with the backbone and applies the projection - head to the output space. If both x0 and x1 are not None, both will be - passed through the backbone and projection head. If x1 is None, only - x0 will be forwarded. - - Args: - x0: - Tensor of shape bsz x channels x W x H. - x1: - Tensor of shape bsz x channels x W x H. - return_features: - Whether or not to return the intermediate features backbone(x). - - Returns: - The output projection of x0 and (if x1 is not None) the output - projection of x1. If return_features is True, the output for each x - is a tuple (out, f) where f are the features before the projection - head. - - Examples: - >>> # single input, single output - >>> out = model(x) - >>> - >>> # single input with return_features=True - >>> out, f = model(x, return_features=True) - >>> - >>> # two inputs, two outputs - >>> out0, out1 = model(x0, x1) - >>> - >>> # two inputs, two outputs with return_features=True - >>> (out0, f0), (out1, f1) = model(x0, x1, return_features=True) - - """ - - # forward pass of first input x0 - f0 = self.backbone(x0).flatten(start_dim=1) - z0 = self.projection_mlp(f0) - p0 = self.prediction_mlp(z0) - - out0 = (z0, p0) - - # append features if requested - if return_features: - out0 = (out0, f0) - - # return out0 if x1 is None - if x1 is None: - return out0 - - # forward pass of second input x1 - f1 = self.backbone(x1).flatten(start_dim=1) - z1 = self.projection_mlp(f1) - p1 = self.prediction_mlp(z1) - - out1 = (z1, p1) - - # append features if requested - if return_features: - out1 = (out1, f1) - - # return both outputs - return out0, out1 diff --git a/lightly/models/simclr.py b/lightly/models/simclr.py deleted file mode 100644 index 905653e09..000000000 --- a/lightly/models/simclr.py +++ /dev/null @@ -1,108 +0,0 @@ -"""SimCLR Model""" - -# Copyright (c) 2020. Lightly AG and its affiliates. -# All Rights Reserved - -import warnings - -import torch -import torch.nn as nn - -from lightly.models.modules import SimCLRProjectionHead - - -class SimCLR(nn.Module): - """Implementation of the SimCLR[0] architecture - - Recommended loss: :py:class:`lightly.loss.ntx_ent_loss.NTXentLoss` - - [0] SimCLR, 2020, https://arxiv.org/abs/2002.05709 - - Attributes: - backbone: - Backbone model to extract features from images. - num_ftrs: - Dimension of the embedding (before the projection head). - out_dim: - Dimension of the output (after the projection head). - - """ - - def __init__(self, backbone: nn.Module, num_ftrs: int = 32, out_dim: int = 128): - super(SimCLR, self).__init__() - - self.backbone = backbone - self.projection_head = SimCLRProjectionHead( - num_ftrs, num_ftrs, out_dim, batch_norm=False - ) - - warnings.warn( - Warning( - "The high-level building block SimCLR will be deprecated in version 1.3.0. " - + "Use low-level building blocks instead. " - + "See https://docs.lightly.ai/self-supervised-learning/lightly.models.html for more information" - ), - DeprecationWarning, - ) - - def forward( - self, x0: torch.Tensor, x1: torch.Tensor = None, return_features: bool = False - ): - """Embeds and projects the input images. - - Extracts features with the backbone and applies the projection - head to the output space. If both x0 and x1 are not None, both will be - passed through the backbone and projection head. If x1 is None, only - x0 will be forwarded. - - Args: - x0: - Tensor of shape bsz x channels x W x H. - x1: - Tensor of shape bsz x channels x W x H. - return_features: - Whether or not to return the intermediate features backbone(x). - - Returns: - The output projection of x0 and (if x1 is not None) the output - projection of x1. If return_features is True, the output for each x - is a tuple (out, f) where f are the features before the projection - head. - - Examples: - >>> # single input, single output - >>> out = model(x) - >>> - >>> # single input with return_features=True - >>> out, f = model(x, return_features=True) - >>> - >>> # two inputs, two outputs - >>> out0, out1 = model(x0, x1) - >>> - >>> # two inputs, two outputs with return_features=True - >>> (out0, f0), (out1, f1) = model(x0, x1, return_features=True) - - """ - - # forward pass of first input x0 - f0 = self.backbone(x0).flatten(start_dim=1) - out0 = self.projection_head(f0) - - # append features if requested - if return_features: - out0 = (out0, f0) - - # return out0 if x1 is None - if x1 is None: - return out0 - - # forward pass of second input x1 - f1 = self.backbone(x1).flatten(start_dim=1) - out1 = self.projection_head(f1) - - # append features if requested - if return_features: - out1 = (out1, f1) - - # return both outputs - return out0, out1 diff --git a/lightly/models/simsiam.py b/lightly/models/simsiam.py deleted file mode 100644 index 2a23d35e2..000000000 --- a/lightly/models/simsiam.py +++ /dev/null @@ -1,134 +0,0 @@ -"""SimSiam Model""" - -# Copyright (c) 2020. Lightly AG and its affiliates. -# All Rights Reserved - -import warnings - -import torch -import torch.nn as nn - -from lightly.models.modules import SimSiamPredictionHead, SimSiamProjectionHead - - -class SimSiam(nn.Module): - """Implementation of SimSiam[0] network - - Recommended loss: :py:class:`lightly.loss.sym_neg_cos_sim_loss.SymNegCosineSimilarityLoss` - - [0] SimSiam, 2020, https://arxiv.org/abs/2011.10566 - - Attributes: - backbone: - Backbone model to extract features from images. - num_ftrs: - Dimension of the embedding (before the projection head). - proj_hidden_dim: - Dimension of the hidden layer of the projection head. This should - be the same size as `num_ftrs`. - pred_hidden_dim: - Dimension of the hidden layer of the predicion head. This should - be `num_ftrs` / 4. - out_dim: - Dimension of the output (after the projection head). - - """ - - def __init__( - self, - backbone: nn.Module, - num_ftrs: int = 2048, - proj_hidden_dim: int = 2048, - pred_hidden_dim: int = 512, - out_dim: int = 2048, - ): - super(SimSiam, self).__init__() - - self.backbone = backbone - self.num_ftrs = num_ftrs - self.proj_hidden_dim = proj_hidden_dim - self.pred_hidden_dim = pred_hidden_dim - self.out_dim = out_dim - - self.projection_mlp = SimSiamProjectionHead( - num_ftrs, - proj_hidden_dim, - out_dim, - ) - - self.prediction_mlp = SimSiamPredictionHead( - out_dim, - pred_hidden_dim, - out_dim, - ) - - warnings.warn( - Warning( - "The high-level building block SimSiam will be deprecated in version 1.3.0. " - + "Use low-level building blocks instead. " - + "See https://docs.lightly.ai/self-supervised-learning/lightly.models.html for more information" - ), - DeprecationWarning, - ) - - def forward( - self, x0: torch.Tensor, x1: torch.Tensor = None, return_features: bool = False - ): - """Forward pass through SimSiam. - - Extracts features with the backbone and applies the projection - head and prediction head to the output space. If both x0 and x1 are not - None, both will be passed through the backbone, projection, and - prediction head. If x1 is None, only x0 will be forwarded. - - Args: - x0: - Tensor of shape bsz x channels x W x H. - x1: - Tensor of shape bsz x channels x W x H. - return_features: - Whether or not to return the intermediate features backbone(x). - - Returns: - The output prediction and projection of x0 and (if x1 is not None) - the output prediction and projection of x1. If return_features is - True, the output for each x is a tuple (out, f) where f are the - features before the projection head. - - Examples: - >>> # single input, single output - >>> out = model(x) - >>> - >>> # single input with return_features=True - >>> out, f = model(x, return_features=True) - >>> - >>> # two inputs, two outputs - >>> out0, out1 = model(x0, x1) - >>> - >>> # two inputs, two outputs with return_features=True - >>> (out0, f0), (out1, f1) = model(x0, x1, return_features=True) - """ - f0 = self.backbone(x0).flatten(start_dim=1) - z0 = self.projection_mlp(f0) - p0 = self.prediction_mlp(z0) - - out0 = (z0, p0) - - # append features if requested - if return_features: - out0 = (out0, f0) - - if x1 is None: - return out0 - - f1 = self.backbone(x1).flatten(start_dim=1) - z1 = self.projection_mlp(f1) - p1 = self.prediction_mlp(z1) - - out1 = (z1, p1) - - # append features if requested - if return_features: - out1 = (out1, f1) - - return out0, out1 diff --git a/lightly/utils/benchmarking/benchmark_module.py b/lightly/utils/benchmarking/benchmark_module.py index 025cd0d25..ba3b81fcc 100644 --- a/lightly/utils/benchmarking/benchmark_module.py +++ b/lightly/utils/benchmarking/benchmark_module.py @@ -53,21 +53,25 @@ class BenchmarkModule(LightningModule): >>> *list(resnet.children())[:-1], >>> nn.AdaptiveAvgPool2d(1), >>> ) - >>> self.resnet_simsiam = - >>> lightly.models.SimSiam(self.backbone, num_ftrs=512) - >>> self.criterion = lightly.loss.SymNegCosineSimilarityLoss() + >>> self.projection_head = SimSiamProjectionHead(512, 512, 128) + >>> self.prediction_head = SimSiamPredictionHead(128, 64, 128) + >>> self.criterion = lightly.loss.NegativeCosineSimilarity() >>> >>> def forward(self, x): - >>> self.resnet_simsiam(x) + >>> f = self.backbone(x).flatten(start_dim=1) + >>> z = self.projection_head(f) + >>> p = self.prediction_head(z) + >>> return z.detach(), p >>> >>> def training_step(self, batch, batch_idx): >>> (x0, x1), _, _ = batch - >>> x0, x1 = self.resnet_simsiam(x0, x1) - >>> loss = self.criterion(x0, x1) + >>> z0, p0 = self.forward(x0) + >>> z1, p1 = self.forward(x1) + >>> loss = 0.5 * (self.criterion(z0, p1) + self.criterion(z1, p0)) >>> return loss >>> def configure_optimizers(self): >>> optim = torch.optim.SGD( - >>> self.resnet_simsiam.parameters(), lr=6e-2, momentum=0.9 + >>> self.parameters(), lr=6e-2, momentum=0.9 >>> ) >>> return [optim] >>> diff --git a/pyproject.toml b/pyproject.toml index 11283481d..a7d4be044 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -127,20 +127,6 @@ known-first-party = ["lightly"] [tool.ruff.lint.pydocstyle] convention = "google" -[tool.ruff.lint.per-file-ignores] -# Deprecated paths (mirrors pyproject.toml's mypy excludes for deprecated code). -"lightly/models/simclr.py" = ["D"] -"lightly/models/moco.py" = ["D"] -"lightly/models/barlowtwins.py" = ["D"] -"lightly/models/nnclr.py" = ["D"] -"lightly/models/simsiam.py" = ["D"] -"lightly/models/byol.py" = ["D"] -"tests/models/test_ModelsSimSiam.py" = ["D"] -"tests/models/test_ModelsSimCLR.py" = ["D"] -"tests/models/test_ModelsNNCLR.py" = ["D"] -"tests/models/test_ModelsMoCo.py" = ["D"] -"tests/models/test_ModelsBYOL.py" = ["D"] - [tool.ruff.format] docstring-code-format = true @@ -223,20 +209,7 @@ exclude = '''(?x)( tests/utils/benchmarking/test_linear_classifier.py | tests/utils/benchmarking/test_metric_callback.py | tests/utils/test_dist.py | - tests/conftest.py | - # Let's not type deprecated models: - lightly/models/simclr.py | - lightly/models/moco.py | - lightly/models/barlowtwins.py | - lightly/models/nnclr.py | - lightly/models/simsiam.py | - lightly/models/byol.py | - # Let's not type deprecated models tests: - tests/models/test_ModelsSimSiam.py | - tests/models/test_ModelsSimCLR.py | - tests/models/test_ModelsNNCLR.py | - tests/models/test_ModelsMoCo.py | - tests/models/test_ModelsBYOL.py )''' + tests/conftest.py )''' # Ignore imports from untyped modules. [[tool.mypy.overrides]] diff --git a/tests/models/test_ModelsBYOL.py b/tests/models/test_ModelsBYOL.py deleted file mode 100644 index 96f2dadb8..000000000 --- a/tests/models/test_ModelsBYOL.py +++ /dev/null @@ -1,112 +0,0 @@ -import pytest -import torch -import torch.nn as nn -import torchvision - -import lightly -from lightly.models import BYOL, ResNetGenerator - - -def get_backbone(resnet, num_ftrs=64): - last_conv_channels = list(resnet.children())[-1].in_features - backbone = nn.Sequential( - lightly.models.batchnorm.get_norm_layer(3, 0), - *list(resnet.children())[:-1], - nn.Conv2d(last_conv_channels, num_ftrs, 1), - nn.AdaptiveAvgPool2d(1), - ) - return backbone - - -class TestModelsBYOL: - @pytest.fixture(autouse=True) - def setup(self): - self.resnet_variants = ["resnet-18", "resnet-50"] - self.batch_size = 2 - self.input_tensor = torch.rand((self.batch_size, 3, 32, 32)) - - def test_create_variations_cpu(self): - for model_name in self.resnet_variants: - resnet = ResNetGenerator(model_name) - model = BYOL(get_backbone(resnet)) - assert model is not None - - def test_create_variations_gpu(self): - device = "cuda" if torch.cuda.is_available() else "cpu" - if device == "cuda": - for model_name in self.resnet_variants: - resnet = ResNetGenerator(model_name) - model = BYOL(get_backbone(resnet)).to(device) - assert model is not None - else: - pass - - def test_feature_dim_configurable(self): - device = "cuda" if torch.cuda.is_available() else "cpu" - for model_name in self.resnet_variants: - for num_ftrs, out_dim in zip([16, 64], [64, 256]): - resnet = ResNetGenerator(model_name) - model = BYOL( - get_backbone(resnet, num_ftrs=num_ftrs), - num_ftrs=num_ftrs, - out_dim=out_dim, - ).to(device) - - # check that feature vector has correct dimension - with torch.no_grad(): - out_features = model.backbone(self.input_tensor.to(device)) - assert out_features.shape[1] == num_ftrs - - # check that projection head output has right dimension - with torch.no_grad(): - out_projection = model.projection_head(out_features.squeeze()) - assert out_projection.shape[1] == out_dim - assert model is not None - - def test_variations_input_dimension(self): - device = "cuda" if torch.cuda.is_available() else "cpu" - for model_name in self.resnet_variants: - for input_width, input_height in zip([32, 64], [64, 64]): - resnet = ResNetGenerator(model_name) - model = BYOL(get_backbone(resnet, num_ftrs=32), num_ftrs=32).to(device) - - input_tensor = torch.rand( - (self.batch_size, 3, input_height, input_width) - ) - with torch.no_grad(): - out, _ = model(input_tensor.to(device), input_tensor.to(device)) - - assert model is not None - assert out is not None - - def test_tuple_input(self): - device = "cuda" if torch.cuda.is_available() else "cpu" - resnet = ResNetGenerator("resnet-18") - model = BYOL(get_backbone(resnet, num_ftrs=32), num_ftrs=32, out_dim=128).to( - device - ) - - x0 = torch.rand((self.batch_size, 3, 64, 64)).to(device) - x1 = torch.rand((self.batch_size, 3, 64, 64)).to(device) - - (z0, p0), (z1, p1) = model(x0, x1) - assert z0.shape == (self.batch_size, 128) - assert z1.shape == (self.batch_size, 128) - assert p0.shape == (self.batch_size, 128) - assert p1.shape == (self.batch_size, 128) - - def test_raises(self): - resnet = ResNetGenerator("resnet-18") - model = BYOL(get_backbone(resnet)) - x0 = torch.rand((self.batch_size, 3, 64, 64)) - - with pytest.raises(ValueError): - model(x0, None) - - with pytest.raises(ValueError): - model(None, x0) - - # test different input shape - x1 = torch.rand((self.batch_size, 5, 32, 32)) - with pytest.raises(ValueError): - model(x0, x1) diff --git a/tests/models/test_ModelsMoCo.py b/tests/models/test_ModelsMoCo.py deleted file mode 100644 index 178b1defd..000000000 --- a/tests/models/test_ModelsMoCo.py +++ /dev/null @@ -1,105 +0,0 @@ -import pytest -import torch -import torch.nn as nn -import torchvision - -import lightly -from lightly.models import MoCo, ResNetGenerator - - -def get_backbone(resnet, num_ftrs=64): - last_conv_channels = list(resnet.children())[-1].in_features - backbone = nn.Sequential( - lightly.models.batchnorm.get_norm_layer(3, 0), - *list(resnet.children())[:-1], - nn.Conv2d(last_conv_channels, num_ftrs, 1), - nn.AdaptiveAvgPool2d(1), - ) - return backbone - - -class TestModelsMoCo: - @pytest.fixture(autouse=True) - def setup(self): - self.resnet_variants = ["resnet-18", "resnet-50"] - self.batch_size = 2 - self.input_tensor = torch.rand((self.batch_size, 3, 32, 32)) - - def test_create_variations_cpu(self): - for model_name in self.resnet_variants: - resnet = ResNetGenerator(model_name) - model = MoCo(get_backbone(resnet)) - assert model is not None - - def test_create_variations_gpu(self): - device = "cuda" if torch.cuda.is_available() else "cpu" - if device == "cuda": - for model_name in self.resnet_variants: - resnet = ResNetGenerator(model_name) - model = MoCo(get_backbone(resnet)).to(device) - assert model is not None - else: - pass - - def test_feature_dim_configurable(self): - device = "cuda" if torch.cuda.is_available() else "cpu" - for model_name in self.resnet_variants: - for num_ftrs, out_dim in zip([16, 64], [64, 256]): - resnet = ResNetGenerator(model_name) - model = MoCo( - get_backbone(resnet, num_ftrs=num_ftrs), - num_ftrs=num_ftrs, - out_dim=out_dim, - ).to(device) - - # check that feature vector has correct dimension - with torch.no_grad(): - out_features = model.backbone(self.input_tensor.to(device)) - assert out_features.shape[1] == num_ftrs - - # check that projection head output has right dimension - with torch.no_grad(): - out_projection = model.projection_head(out_features.squeeze()) - assert out_projection.shape[1] == out_dim - assert model is not None - - def test_variations_input_dimension(self): - device = "cuda" if torch.cuda.is_available() else "cpu" - for model_name in self.resnet_variants: - for input_width, input_height in zip([32, 64], [64, 64]): - resnet = ResNetGenerator(model_name) - model = MoCo(get_backbone(resnet, num_ftrs=32)).to(device) - - input_tensor = torch.rand( - (self.batch_size, 3, input_height, input_width) - ) - with torch.no_grad(): - out = model(input_tensor.to(device)) - - assert model is not None - assert out is not None - - def test_tuple_input(self): - device = "cuda" if torch.cuda.is_available() else "cpu" - resnet = ResNetGenerator("resnet-18") - model = MoCo(get_backbone(resnet, num_ftrs=32), out_dim=128).to(device) - - x0 = torch.rand((self.batch_size, 3, 64, 64)).to(device) - x1 = torch.rand((self.batch_size, 3, 64, 64)).to(device) - - out = model(x0) - assert out.shape == (self.batch_size, 128) - - out, features = model(x0, return_features=True) - assert out.shape == (self.batch_size, 128) - assert features.shape == (self.batch_size, 32) - - out0, out1 = model(x0, x1) - assert out0.shape == (self.batch_size, 128) - assert out1.shape == (self.batch_size, 128) - - (out0, f0), (out1, f1) = model(x0, x1, return_features=True) - assert out0.shape == (self.batch_size, 128) - assert out1.shape == (self.batch_size, 128) - assert f0.shape == (self.batch_size, 32) - assert f1.shape == (self.batch_size, 32) diff --git a/tests/models/test_ModelsNNCLR.py b/tests/models/test_ModelsNNCLR.py deleted file mode 100644 index 64f0aaf84..000000000 --- a/tests/models/test_ModelsNNCLR.py +++ /dev/null @@ -1,126 +0,0 @@ -import pytest -import torch -import torch.nn as nn -import torchvision - -from lightly.models import NNCLR -from lightly.models.modules import NNMemoryBankModule - - -def resnet_generator(name: str): - if name == "resnet18": - return torchvision.models.resnet18() - elif name == "resnet50": - return torchvision.models.resnet50() - raise NotImplementedError - - -def get_backbone(model: nn.Module): - backbone = torch.nn.Sequential(*(list(model.children())[:-1])) - return backbone - - -class TestNNCLR: - @pytest.fixture(autouse=True) - def setup(self): - self.resnet_variants = dict( - resnet18=dict( - num_ftrs=512, - proj_hidden_dim=512, - pred_hidden_dim=128, - out_dim=512, - ), - resnet50=dict( - num_ftrs=2048, - proj_hidden_dim=2048, - pred_hidden_dim=512, - out_dim=2048, - ), - ) - self.batch_size = 2 - self.input_tensor = torch.rand((self.batch_size, 3, 32, 32)) - - def test_create_variations_cpu(self): - for model_name, config in self.resnet_variants.items(): - resnet = resnet_generator(model_name) - model = NNCLR(get_backbone(resnet), **config) - assert model is not None - - def test_create_variations_gpu(self): - if not torch.cuda.is_available(): - return - - for model_name, config in self.resnet_variants.items(): - resnet = resnet_generator(model_name) - model = NNCLR(get_backbone(resnet), **config).to("cuda") - assert model is not None - - def test_feature_dim_configurable(self): - device = "cuda" if torch.cuda.is_available() else "cpu" - for model_name, config in self.resnet_variants.items(): - resnet = resnet_generator(model_name) - model = NNCLR(get_backbone(resnet), **config).to(device) - - # check that feature vector has correct dimension - with torch.no_grad(): - out_features = model.backbone(self.input_tensor.to(device)) - assert out_features.shape[1] == config["num_ftrs"] - - # check that projection head output has right dimension - with torch.no_grad(): - out_projection = model.projection_mlp(out_features.squeeze()) - assert out_projection.shape[1] == config["out_dim"] - - # check that prediction head output has right dimension - with torch.no_grad(): - out_prediction = model.prediction_mlp(out_projection.squeeze()) - assert out_prediction.shape[1] == config["out_dim"] - - def test_tuple_input(self): - device = "cuda" if torch.cuda.is_available() else "cpu" - for model_name, config in self.resnet_variants.items(): - resnet = resnet_generator(model_name) - model = NNCLR(get_backbone(resnet), **config).to(device) - - x0 = torch.rand((self.batch_size, 3, 64, 64)).to(device) - x1 = torch.rand((self.batch_size, 3, 64, 64)).to(device) - - out = model(x0) - assert out[0].shape == (self.batch_size, config["out_dim"]) - assert out[1].shape == (self.batch_size, config["out_dim"]) - - out, features = model(x0, return_features=True) - assert out[0].shape == (self.batch_size, config["out_dim"]) - assert out[1].shape == (self.batch_size, config["out_dim"]) - assert features.shape == (self.batch_size, config["num_ftrs"]) - - out0, out1 = model(x0, x1) - assert out0[0].shape == (self.batch_size, config["out_dim"]) - assert out0[1].shape == (self.batch_size, config["out_dim"]) - assert out1[0].shape == (self.batch_size, config["out_dim"]) - assert out1[1].shape == (self.batch_size, config["out_dim"]) - - (out0, f0), (out1, f1) = model(x0, x1, return_features=True) - assert out0[0].shape == (self.batch_size, config["out_dim"]) - assert out0[1].shape == (self.batch_size, config["out_dim"]) - assert out1[0].shape == (self.batch_size, config["out_dim"]) - assert out1[1].shape == (self.batch_size, config["out_dim"]) - assert f0.shape == (self.batch_size, config["num_ftrs"]) - assert f1.shape == (self.batch_size, config["num_ftrs"]) - - def test_memory_bank(self): - device = "cuda" if torch.cuda.is_available() else "cpu" - for model_name, config in self.resnet_variants.items(): - resnet = resnet_generator(model_name) - model = NNCLR(get_backbone(resnet), **config).to(device) - - for nn_size in [2**3, 2**8]: - nn_replacer = NNMemoryBankModule(size=(nn_size, config["out_dim"])) - - with torch.no_grad(): - for i in range(10): - x0 = torch.rand((self.batch_size, 3, 64, 64)).to(device) - x1 = torch.rand((self.batch_size, 3, 64, 64)).to(device) - (z0, p0), (z1, p1) = model(x0, x1) - z0 = nn_replacer(z0.detach(), update=False) - z1 = nn_replacer(z1.detach(), update=True) diff --git a/tests/models/test_ModelsSimCLR.py b/tests/models/test_ModelsSimCLR.py deleted file mode 100644 index e19c3a42e..000000000 --- a/tests/models/test_ModelsSimCLR.py +++ /dev/null @@ -1,105 +0,0 @@ -import pytest -import torch -import torch.nn as nn -import torchvision - -import lightly -from lightly.models import ResNetGenerator, SimCLR - - -def get_backbone(resnet, num_ftrs=64): - last_conv_channels = list(resnet.children())[-1].in_features - backbone = nn.Sequential( - lightly.models.batchnorm.get_norm_layer(3, 0), - *list(resnet.children())[:-1], - nn.Conv2d(last_conv_channels, num_ftrs, 1), - nn.AdaptiveAvgPool2d(1), - ) - return backbone - - -class TestModelsSimCLR: - @pytest.fixture(autouse=True) - def setup(self): - self.resnet_variants = ["resnet-18", "resnet-50"] - self.batch_size = 2 - self.input_tensor = torch.rand((self.batch_size, 3, 32, 32)) - - def test_create_variations_cpu(self): - for model_name in self.resnet_variants: - resnet = ResNetGenerator(model_name) - model = SimCLR(get_backbone(resnet)) - assert model is not None - - def test_create_variations_gpu(self): - device = "cuda" if torch.cuda.is_available() else "cpu" - if device == "cuda": - for model_name in self.resnet_variants: - resnet = ResNetGenerator(model_name) - model = SimCLR(get_backbone(resnet)).to(device) - assert model is not None - else: - pass - - def test_feature_dim_configurable(self): - device = "cuda" if torch.cuda.is_available() else "cpu" - for model_name in self.resnet_variants: - for num_ftrs, out_dim in zip([16, 64], [64, 256]): - resnet = ResNetGenerator(model_name) - model = SimCLR( - get_backbone(resnet, num_ftrs=num_ftrs), - num_ftrs=num_ftrs, - out_dim=out_dim, - ).to(device) - - # check that feature vector has correct dimension - with torch.no_grad(): - out_features = model.backbone(self.input_tensor.to(device)) - assert out_features.shape[1] == num_ftrs - - # check that projection head output has right dimension - with torch.no_grad(): - out_projection = model.projection_head(out_features.squeeze()) - assert out_projection.shape[1] == out_dim - assert model is not None - - def test_variations_input_dimension(self): - device = "cuda" if torch.cuda.is_available() else "cpu" - for model_name in self.resnet_variants: - for input_width, input_height in zip([32, 64], [64, 64]): - resnet = ResNetGenerator(model_name) - model = SimCLR(get_backbone(resnet, num_ftrs=32)).to(device) - - input_tensor = torch.rand( - (self.batch_size, 3, input_height, input_width) - ) - with torch.no_grad(): - out = model(input_tensor.to(device)) - - assert model is not None - assert out is not None - - def test_tuple_input(self): - device = "cuda" if torch.cuda.is_available() else "cpu" - resnet = ResNetGenerator("resnet-18") - model = SimCLR(get_backbone(resnet, num_ftrs=32), out_dim=128).to(device) - - x0 = torch.rand((self.batch_size, 3, 64, 64)).to(device) - x1 = torch.rand((self.batch_size, 3, 64, 64)).to(device) - - out = model(x0) - assert out.shape == (self.batch_size, 128) - - out, features = model(x0, return_features=True) - assert out.shape == (self.batch_size, 128) - assert features.shape == (self.batch_size, 32) - - out0, out1 = model(x0, x1) - assert out0.shape == (self.batch_size, 128) - assert out1.shape == (self.batch_size, 128) - - (out0, f0), (out1, f1) = model(x0, x1, return_features=True) - assert out0.shape == (self.batch_size, 128) - assert out1.shape == (self.batch_size, 128) - assert f0.shape == (self.batch_size, 32) - assert f1.shape == (self.batch_size, 32) diff --git a/tests/models/test_ModelsSimSiam.py b/tests/models/test_ModelsSimSiam.py deleted file mode 100644 index 18a53b3d9..000000000 --- a/tests/models/test_ModelsSimSiam.py +++ /dev/null @@ -1,108 +0,0 @@ -import pytest -import torch -import torch.nn as nn -import torchvision - -from lightly.models import SimSiam - - -def resnet_generator(name: str): - if name == "resnet18": - return torchvision.models.resnet18() - elif name == "resnet50": - return torchvision.models.resnet50() - raise NotImplementedError - - -def get_backbone(model: nn.Module): - backbone = torch.nn.Sequential(*(list(model.children())[:-1])) - return backbone - - -class TestSimSiam: - @pytest.fixture(autouse=True) - def setup(self): - self.resnet_variants = dict( - resnet18=dict( - num_ftrs=512, - proj_hidden_dim=512, - pred_hidden_dim=128, - out_dim=512, - ), - resnet50=dict( - num_ftrs=2048, - proj_hidden_dim=2048, - pred_hidden_dim=512, - out_dim=2048, - ), - ) - self.batch_size = 2 - self.input_tensor = torch.rand((self.batch_size, 3, 32, 32)) - - def test_create_variations_cpu(self): - for model_name, config in self.resnet_variants.items(): - resnet = resnet_generator(model_name) - model = SimSiam(get_backbone(resnet), **config) - assert model is not None - - def test_create_variations_gpu(self): - if not torch.cuda.is_available(): - return - - for model_name, config in self.resnet_variants.items(): - resnet = resnet_generator(model_name) - model = SimSiam(get_backbone(resnet), **config).to("cuda") - assert model is not None - - def test_feature_dim_configurable(self): - device = "cuda" if torch.cuda.is_available() else "cpu" - for model_name, config in self.resnet_variants.items(): - resnet = resnet_generator(model_name) - model = SimSiam(get_backbone(resnet), **config).to(device) - - # check that feature vector has correct dimension - with torch.no_grad(): - out_features = model.backbone(self.input_tensor.to(device)) - assert out_features.shape[1] == config["num_ftrs"] - - # check that projection head output has right dimension - with torch.no_grad(): - out_projection = model.projection_mlp(out_features.squeeze()) - assert out_projection.shape[1] == config["out_dim"] - - # check that prediction head output has right dimension - with torch.no_grad(): - out_prediction = model.prediction_mlp(out_projection.squeeze()) - assert out_prediction.shape[1] == config["out_dim"] - - def test_tuple_input(self): - device = "cuda" if torch.cuda.is_available() else "cpu" - for model_name, config in self.resnet_variants.items(): - resnet = resnet_generator(model_name) - model = SimSiam(get_backbone(resnet), **config).to(device) - - x0 = torch.rand((self.batch_size, 3, 64, 64)).to(device) - x1 = torch.rand((self.batch_size, 3, 64, 64)).to(device) - - out = model(x0) - assert out[0].shape == (self.batch_size, config["out_dim"]) - assert out[1].shape == (self.batch_size, config["out_dim"]) - - out, features = model(x0, return_features=True) - assert out[0].shape == (self.batch_size, config["out_dim"]) - assert out[1].shape == (self.batch_size, config["out_dim"]) - assert features.shape == (self.batch_size, config["num_ftrs"]) - - out0, out1 = model(x0, x1) - assert out0[0].shape == (self.batch_size, config["out_dim"]) - assert out0[1].shape == (self.batch_size, config["out_dim"]) - assert out1[0].shape == (self.batch_size, config["out_dim"]) - assert out1[1].shape == (self.batch_size, config["out_dim"]) - - (out0, f0), (out1, f1) = model(x0, x1, return_features=True) - assert out0[0].shape == (self.batch_size, config["out_dim"]) - assert out0[1].shape == (self.batch_size, config["out_dim"]) - assert out1[0].shape == (self.batch_size, config["out_dim"]) - assert out1[1].shape == (self.batch_size, config["out_dim"]) - assert f0.shape == (self.batch_size, config["num_ftrs"]) - assert f1.shape == (self.batch_size, config["num_ftrs"]) From 07221cdcd7b92c3c6f089ad2d6b98d7e244dda2d Mon Sep 17 00:00:00 2001 From: gabrielfruet Date: Tue, 11 Aug 2026 18:10:48 -0300 Subject: [PATCH 3/9] Remove deprecated GaussianBlur kernel_size and scale Both arguments have warned since 1.4.0 and were ignored after warning, so removing them changes no augmentation output. The pass-through arguments in the method transforms go with them. FDATransform defaulted kernel_size to 23 rather than None, so constructing one emitted two DeprecationWarnings. That is gone now. GaussianBlur.__init__ is keyword-only. kernel_size was the first positional argument, so dropping it would otherwise reinterpret GaussianBlur(0.3) as prob=0.3. detcon_transform keeps its kernel_size: it feeds torchvision's GaussianBlur, which is a different argument. Co-Authored-By: Claude Opus 5 --- lightly/cli/config/config.yaml | 3 --- lightly/transforms/byol_transform.py | 16 ++++--------- lightly/transforms/densecl_transform.py | 4 ---- lightly/transforms/dino_transform.py | 27 +-------------------- lightly/transforms/fast_siam_transform.py | 6 ----- lightly/transforms/fda_transform.py | 16 ++++--------- lightly/transforms/gaussian_blur.py | 19 ++------------- lightly/transforms/ibot_transform.py | 29 ++--------------------- lightly/transforms/mmcr_transform.py | 2 -- lightly/transforms/moco_transform.py | 12 ---------- lightly/transforms/msn_transform.py | 17 +++---------- lightly/transforms/simclr_transform.py | 15 +++--------- lightly/transforms/simsiam_transform.py | 14 ++--------- lightly/transforms/smog_transform.py | 16 ++----------- lightly/transforms/swav_transform.py | 14 ++--------- lightly/transforms/vicreg_transform.py | 14 ++--------- lightly/transforms/vicregl_transform.py | 29 ++++------------------- lightly/transforms/wmse_transform.py | 11 ++------- tests/transforms/test_gaussian_blur.py | 10 +++----- 19 files changed, 36 insertions(+), 238 deletions(-) diff --git a/lightly/cli/config/config.yaml b/lightly/cli/config/config.yaml index 1a5c6bcad..40e879ada 100644 --- a/lightly/cli/config/config.yaml +++ b/lightly/cli/config/config.yaml @@ -44,9 +44,6 @@ collate: random_gray_scale: 0.2 # Probability of converting image to gray scale. gaussian_blur: 0.5 # Probability of Gaussian blur. sigmas: [0.2, 2] # Sigmas of Gaussian blur - kernel_size: null # Will be deprecated in favor of `sigmas` argument. If set, the old behavior - # applies and `sigmas` is ignored. Used to calculate sigma of gaussian blur - # with kernel_size * input_size. vf_prob: 0.0 # Probability that vertical flip is applied. hf_prob: 0.5 # Probability that horizontal flip is applied. rr_prob: 0.0 # Probability that random rotation is applied. diff --git a/lightly/transforms/byol_transform.py b/lightly/transforms/byol_transform.py index 6a30e857e..cc06d7a91 100644 --- a/lightly/transforms/byol_transform.py +++ b/lightly/transforms/byol_transform.py @@ -46,7 +46,6 @@ def __init__( random_gray_scale: float = 0.2, gaussian_blur: float = 1.0, solarization_prob: float = 0.0, - kernel_size: Optional[float] = None, sigmas: Tuple[float, float] = (0.1, 2), vf_prob: float = 0.0, hf_prob: float = 0.5, @@ -71,11 +70,8 @@ def __init__( random_gray_scale: Probability of conversion to grayscale. gaussian_blur: Probability of Gaussian blur. solarization_prob: Probability of solarization. - kernel_size: Will be deprecated in favor of `sigmas` argument. If set, - the old behavior applies and `sigmas` is ignored. Used to calculate - sigma of gaussian blur with kernel_size * input_size. sigmas: Tuple of min and max value from which the std of the gaussian - kernel is sampled. Is ignored if `kernel_size` is set. + kernel is sampled. vf_prob: Probability that vertical flip is applied. hf_prob: Probability that horizontal flip is applied. rr_prob: Probability that random rotation is applied. @@ -103,7 +99,7 @@ def __init__( random_rotation_transform(rr_prob=rr_prob, rr_degrees=rr_degrees), T.RandomApply([color_jitter], p=cj_prob), T.RandomGrayscale(p=random_gray_scale), - GaussianBlur(kernel_size=kernel_size, sigmas=sigmas, prob=gaussian_blur), + GaussianBlur(sigmas=sigmas, prob=gaussian_blur), RandomSolarization(prob=solarization_prob), T.ToTensor(), ] @@ -161,7 +157,6 @@ def __init__( random_gray_scale: float = 0.2, gaussian_blur: float = 0.1, solarization_prob: float = 0.2, - kernel_size: Optional[float] = None, sigmas: Tuple[float, float] = (0.1, 2), vf_prob: float = 0.0, hf_prob: float = 0.5, @@ -186,11 +181,8 @@ def __init__( random_gray_scale: Probability of conversion to grayscale. gaussian_blur: Probability of Gaussian blur. solarization_prob: Probability of solarization. - kernel_size: Will be deprecated in favor of `sigmas` argument. If set, - the old behavior applies and `sigmas` is ignored. Used to calculate - sigma of gaussian blur with kernel_size * input_size. sigmas: Tuple of min and max value from which the std of the gaussian - kernel is sampled. Is ignored if `kernel_size` is set. + kernel is sampled. vf_prob: Probability that vertical flip is applied. hf_prob: Probability that horizontal flip is applied. rr_prob: Probability that random rotation is applied. @@ -218,7 +210,7 @@ def __init__( random_rotation_transform(rr_prob=rr_prob, rr_degrees=rr_degrees), T.RandomApply([color_jitter], p=cj_prob), T.RandomGrayscale(p=random_gray_scale), - GaussianBlur(kernel_size=kernel_size, sigmas=sigmas, prob=gaussian_blur), + GaussianBlur(sigmas=sigmas, prob=gaussian_blur), RandomSolarization(prob=solarization_prob), T.ToTensor(), ] diff --git a/lightly/transforms/densecl_transform.py b/lightly/transforms/densecl_transform.py index 45806687c..da5547274 100644 --- a/lightly/transforms/densecl_transform.py +++ b/lightly/transforms/densecl_transform.py @@ -45,12 +45,8 @@ class DenseCLTransform(MoCoV2Transform): Probability of conversion to grayscale. gaussian_blur: Probability of Gaussian blur. - kernel_size: - Will be deprecated in favor of `sigmas` argument. If set, the old behavior applies and `sigmas` is ignored. - Used to calculate sigma of gaussian blur with kernel_size * input_size. sigmas: Tuple of min and max value from which the std of the gaussian kernel is sampled. - Is ignored if `kernel_size` is set. vf_prob: Probability that vertical flip is applied. hf_prob: diff --git a/lightly/transforms/dino_transform.py b/lightly/transforms/dino_transform.py index 5a897ed74..c92433848 100644 --- a/lightly/transforms/dino_transform.py +++ b/lightly/transforms/dino_transform.py @@ -77,15 +77,8 @@ class DINOTransform(MultiViewTransform): Tuple of probabilities to apply gaussian blur on the different views. The input is ordered as follows: (global_view_0, global_view_1, local_views) - kernel_size: - Will be deprecated in favor of `sigmas` argument. If set, the old behavior applies and `sigmas` is ignored. - Used to calculate sigma of gaussian blur with kernel_size * input_size. - kernel_scale: - Old argument. Value is deprecated in favor of sigmas. If set, the old behavior applies and `sigmas` is ignored. - Used to scale the `kernel_size` of a factor of `kernel_scale` sigmas: Tuple of min and max value from which the std of the gaussian kernel is sampled. - Is ignored if `kernel_size` is set. solarization: Probability to apply solarization on the second global view. normalize: @@ -112,8 +105,6 @@ def __init__( cj_hue: float = 0.2, random_gray_scale: float = 0.2, gaussian_blur: Tuple[float, float, float] = (1.0, 0.1, 0.5), - kernel_size: Optional[float] = None, - kernel_scale: Optional[float] = None, sigmas: Tuple[float, float] = (0.1, 2), solarization_prob: float = 0.2, normalize: Union[None, Dict[str, List[float]]] = IMAGENET_NORMALIZE, @@ -134,8 +125,6 @@ def __init__( cj_sat=cj_sat, random_gray_scale=random_gray_scale, gaussian_blur=gaussian_blur[0], - kernel_size=kernel_size, - kernel_scale=kernel_scale, sigmas=sigmas, solarization_prob=0, normalize=normalize, @@ -156,8 +145,6 @@ def __init__( cj_sat=cj_sat, random_gray_scale=random_gray_scale, gaussian_blur=gaussian_blur[1], - kernel_size=kernel_size, - kernel_scale=kernel_scale, sigmas=sigmas, solarization_prob=solarization_prob, normalize=normalize, @@ -179,8 +166,6 @@ def __init__( cj_sat=cj_sat, random_gray_scale=random_gray_scale, gaussian_blur=gaussian_blur[2], - kernel_size=kernel_size, - kernel_scale=kernel_scale, sigmas=sigmas, solarization_prob=0, normalize=normalize, @@ -232,8 +217,6 @@ def __init__( cj_hue: float = 0.2, random_gray_scale: float = 0.2, gaussian_blur: float = 1.0, - kernel_size: Optional[float] = None, - kernel_scale: Optional[float] = None, sigmas: Tuple[float, float] = (0.1, 2), solarization_prob: float = 0.2, normalize: Union[None, Dict[str, List[float]]] = IMAGENET_NORMALIZE, @@ -261,14 +244,8 @@ def __init__( cj_hue: How much to jitter hue. random_gray_scale: Probability of conversion to grayscale. gaussian_blur: Probability of Gaussian blur. - kernel_size: Will be deprecated in favor of `sigmas` argument. If set, - the old behavior applies and `sigmas` is ignored. Used to calculate - sigma of gaussian blur with kernel_size * input_size. - kernel_scale: Old argument. Will be deprecated in favor of `sigmas` - argument. If set, the old behavior applies and `sigmas` is ignored. - Used to scale the `kernel_size` of a factor of `kernel_scale`. sigmas: Tuple of min and max value from which the std of the gaussian - kernel is sampled. Is ignored if `kernel_size` is set. + kernel is sampled. solarization_prob: Probability to apply solarization. normalize: Dictionary with 'mean' and 'std' for torchvision.transforms.Normalize. @@ -296,8 +273,6 @@ def __init__( ), T.RandomGrayscale(p=random_gray_scale), GaussianBlur( - kernel_size=kernel_size, - scale=kernel_scale, sigmas=sigmas, prob=gaussian_blur, ), diff --git a/lightly/transforms/fast_siam_transform.py b/lightly/transforms/fast_siam_transform.py index 73d328277..a2840b3c0 100644 --- a/lightly/transforms/fast_siam_transform.py +++ b/lightly/transforms/fast_siam_transform.py @@ -47,12 +47,8 @@ class FastSiamTransform(MultiViewTransform): Probability of conversion to grayscale. gaussian_blur: Probability of Gaussian blur. - kernel_size: - Will be deprecated in favor of `sigmas` argument. If set, the old behavior applies and `sigmas` is ignored. - Used to calculate sigma of gaussian blur with kernel_size * input_size. sigmas: Tuple of min and max value from which the std of the gaussian kernel is sampled. - Is ignored if `kernel_size` is set. vf_prob: Probability that vertical flip is applied. hf_prob: @@ -83,7 +79,6 @@ def __init__( min_scale: float = 0.2, random_gray_scale: float = 0.2, gaussian_blur: float = 0.5, - kernel_size: Optional[float] = None, sigmas: Tuple[float, float] = (0.1, 2), vf_prob: float = 0.0, hf_prob: float = 0.5, @@ -103,7 +98,6 @@ def __init__( min_scale=min_scale, random_gray_scale=random_gray_scale, gaussian_blur=gaussian_blur, - kernel_size=kernel_size, sigmas=sigmas, vf_prob=vf_prob, hf_prob=hf_prob, diff --git a/lightly/transforms/fda_transform.py b/lightly/transforms/fda_transform.py index ff53dde49..b918e3427 100644 --- a/lightly/transforms/fda_transform.py +++ b/lightly/transforms/fda_transform.py @@ -66,7 +66,6 @@ def __init__( # Gaussian blur gaussian_blur: float = 1.0, sigmas: Tuple[float, float] = (0.1, 2), - kernel_size: Optional[float] = 23, # Amplitude rescale ampl_rescale_range: Tuple[float, float] = (0.8, 1.75), ampl_rescale_prob: float = 0.2, @@ -103,10 +102,7 @@ def __init__( random_gray_scale: Probability of conversion to grayscale. gaussian_blur: Probability of Gaussian blur. sigmas: Tuple of min and max value from which the std of the gaussian - kernel is sampled. Is ignored if `kernel_size` is set. - kernel_size: Will be deprecated in favor of `sigmas` argument. If set, - the old behavior applies and `sigmas` is ignored. Used to calculate - sigma of gaussian blur with kernel_size * input_size. + kernel is sampled. ampl_rescale_range: Range of the amplitude rescaling factor for frequency domain augmentation. ampl_rescale_prob: Probability of applying amplitude rescaling. @@ -170,7 +166,7 @@ def __init__( random_rotation_transform(rr_prob=rr_prob, rr_degrees=rr_degrees), T.RandomApply([color_jitter], p=cj_prob), T.RandomGrayscale(p=random_gray_scale), - GaussianBlur(kernel_size=kernel_size, sigmas=sigmas, prob=gaussian_blur), + GaussianBlur(sigmas=sigmas, prob=gaussian_blur), RandomSolarization(prob=solarization_prob), T.ToTensor(), ] @@ -239,7 +235,6 @@ def __init__( # Gaussian blur gaussian_blur: float = 0.1, sigmas: Tuple[float, float] = (0.1, 2), - kernel_size: Optional[float] = 23, # Amplitude rescale ampl_rescale_range: Tuple[float, float] = (0.8, 1.75), ampl_rescale_prob: float = 0.2, @@ -276,10 +271,7 @@ def __init__( random_gray_scale: Probability of conversion to grayscale. gaussian_blur: Probability of Gaussian blur. sigmas: Tuple of min and max value from which the std of the gaussian - kernel is sampled. Is ignored if `kernel_size` is set. - kernel_size: Will be deprecated in favor of `sigmas` argument. If set, - the old behavior applies and `sigmas` is ignored. Used to calculate - sigma of gaussian blur with kernel_size * input_size. + kernel is sampled. ampl_rescale_range: Range of the amplitude rescaling factor for frequency domain augmentation. ampl_rescale_prob: Probability of applying amplitude rescaling. @@ -343,7 +335,7 @@ def __init__( random_rotation_transform(rr_prob=rr_prob, rr_degrees=rr_degrees), T.RandomApply([color_jitter], p=cj_prob), T.RandomGrayscale(p=random_gray_scale), - GaussianBlur(kernel_size=kernel_size, sigmas=sigmas, prob=gaussian_blur), + GaussianBlur(sigmas=sigmas, prob=gaussian_blur), RandomSolarization(prob=solarization_prob), T.ToTensor(), ] diff --git a/lightly/transforms/gaussian_blur.py b/lightly/transforms/gaussian_blur.py index 207f3ab87..d97e54dd3 100644 --- a/lightly/transforms/gaussian_blur.py +++ b/lightly/transforms/gaussian_blur.py @@ -1,8 +1,7 @@ # Copyright (c) 2020. Lightly AG and its affiliates. # All Rights Reserved -from typing import Optional, Tuple, Union -from warnings import warn +from typing import Tuple, Union import numpy as np from PIL import ImageFilter @@ -21,33 +20,19 @@ class GaussianBlur: the Gaussian kernel. Attributes: - kernel_size: - Will be deprecated in favor of `sigmas` argument. If set, the old behavior applies and `sigmas` is ignored. - Used to calculate sigma of gaussian blur with kernel_size * input_size. prob: Probability with which the blur is applied. - scale: - Will be deprecated in favor of `sigmas` argument. If set, the old behavior applies and `sigmas` is ignored. - Used to scale the `kernel_size` of a factor of `kernel_scale` sigmas: Tuple of min and max value from which the std of the gaussian kernel is sampled. - Is ignored if `kernel_size` is set. """ def __init__( self, - kernel_size: Optional[float] = None, + *, prob: float = 0.5, - scale: Optional[float] = None, sigmas: Tuple[float, float] = (0.2, 2), ): - if scale != None or kernel_size != None: - warn( - "The 'kernel_size' and 'scale' arguments of the GaussianBlur augmentation will be deprecated. " - "Please use the 'sigmas' parameter instead.", - DeprecationWarning, - ) self.prob = prob self.sigmas = sigmas diff --git a/lightly/transforms/ibot_transform.py b/lightly/transforms/ibot_transform.py index be8c44fb1..5b76f998b 100644 --- a/lightly/transforms/ibot_transform.py +++ b/lightly/transforms/ibot_transform.py @@ -1,4 +1,4 @@ -from typing import Dict, List, Optional, Tuple, Union +from typing import Dict, List, Tuple, Union from PIL.Image import Image from torch import Tensor @@ -66,15 +66,8 @@ class IBOTTransform(MultiViewTransform): Tuple of probabilities to apply gaussian blur on the different views. The input is ordered as follows: (global_view_0, global_view_1, local_views) - kernel_size: - Will be deprecated in favor of `sigmas` argument. If set, the old behavior applies and `sigmas` is ignored. - Used to calculate sigma of gaussian blur with kernel_size * input_size. - kernel_scale: - Old argument. Value is deprecated in favor of sigmas. If set, the old behavior applies and `sigmas` is ignored. - Used to scale the `kernel_size` of a factor of `kernel_scale` sigmas: Tuple of min and max value from which the std of the gaussian kernel is sampled. - Is ignored if `kernel_size` is set. solarization: Probability to apply solarization on the second global view. normalize: @@ -98,8 +91,6 @@ def __init__( cj_hue: float = 0.2, random_gray_scale: float = 0.2, gaussian_blur: Tuple[float, float, float] = (1.0, 0.1, 0.5), - kernel_size: Optional[float] = None, - kernel_scale: Optional[float] = None, sigmas: Tuple[float, float] = (0.1, 2), solarization_prob: float = 0.2, normalize: Union[None, Dict[str, List[float]]] = IMAGENET_NORMALIZE, @@ -117,8 +108,6 @@ def __init__( cj_sat=cj_sat, random_gray_scale=random_gray_scale, gaussian_blur=gaussian_blur[0], - kernel_size=kernel_size, - kernel_scale=kernel_scale, sigmas=sigmas, solarization_prob=0, normalize=normalize, @@ -136,8 +125,6 @@ def __init__( cj_sat=cj_sat, random_gray_scale=random_gray_scale, gaussian_blur=gaussian_blur[1], - kernel_size=kernel_size, - kernel_scale=kernel_scale, sigmas=sigmas, solarization_prob=solarization_prob, normalize=normalize, @@ -156,8 +143,6 @@ def __init__( cj_sat=cj_sat, random_gray_scale=random_gray_scale, gaussian_blur=gaussian_blur[2], - kernel_size=kernel_size, - kernel_scale=kernel_scale, sigmas=sigmas, solarization_prob=0, normalize=normalize, @@ -204,8 +189,6 @@ def __init__( cj_hue: float = 0.2, random_gray_scale: float = 0.2, gaussian_blur: float = 1.0, - kernel_size: Optional[float] = None, - kernel_scale: Optional[float] = None, sigmas: Tuple[float, float] = (0.1, 2), solarization_prob: float = 0.2, normalize: Union[None, Dict[str, List[float]]] = IMAGENET_NORMALIZE, @@ -225,14 +208,8 @@ def __init__( cj_hue: How much to jitter hue. random_gray_scale: Probability of conversion to grayscale. gaussian_blur: Probability of Gaussian blur. - kernel_size: Will be deprecated in favor of `sigmas` argument. If set, - the old behavior applies and `sigmas` is ignored. Used to calculate - sigma of gaussian blur with kernel_size * input_size. - kernel_scale: Old argument. Will be deprecated in favor of `sigmas` - argument. If set, the old behavior applies and `sigmas` is ignored. - Used to scale the `kernel_size` of a factor of `kernel_scale`. sigmas: Tuple of min and max value from which the std of the gaussian - kernel is sampled. Is ignored if `kernel_size` is set. + kernel is sampled. solarization_prob: Probability to apply solarization. normalize: Dictionary with 'mean' and 'std' for torchvision.transforms.Normalize. @@ -258,8 +235,6 @@ def __init__( ), T.RandomGrayscale(p=random_gray_scale), GaussianBlur( - kernel_size=kernel_size, - scale=kernel_scale, sigmas=sigmas, prob=gaussian_blur, ), diff --git a/lightly/transforms/mmcr_transform.py b/lightly/transforms/mmcr_transform.py index de156e7fb..c4e12e80b 100644 --- a/lightly/transforms/mmcr_transform.py +++ b/lightly/transforms/mmcr_transform.py @@ -56,7 +56,6 @@ def __init__( random_gray_scale: float = 0.2, gaussian_blur: float = 1.0, solarization_prob: float = 0.0, - kernel_size: Optional[float] = None, sigmas: Tuple[float, float] = (0.1, 2), vf_prob: float = 0.0, hf_prob: float = 0.5, @@ -78,7 +77,6 @@ def __init__( random_gray_scale=random_gray_scale, gaussian_blur=gaussian_blur, solarization_prob=solarization_prob, - kernel_size=kernel_size, sigmas=sigmas, vf_prob=vf_prob, hf_prob=hf_prob, diff --git a/lightly/transforms/moco_transform.py b/lightly/transforms/moco_transform.py index e2744b42d..b3b09a5b0 100644 --- a/lightly/transforms/moco_transform.py +++ b/lightly/transforms/moco_transform.py @@ -42,12 +42,8 @@ class MoCoV1Transform(SimCLRTransform): Probability of conversion to grayscale. gaussian_blur: Probability of Gaussian blur. - kernel_size: - Will be deprecated in favor of `sigmas` argument. If set, the old behavior applies and `sigmas` is ignored. - Used to calculate sigma of gaussian blur with kernel_size * input_size. sigmas: Tuple of min and max value from which the std of the gaussian kernel is sampled. - Is ignored if `kernel_size` is set. vf_prob: Probability that vertical flip is applied. hf_prob: @@ -77,7 +73,6 @@ def __init__( min_scale: float = 0.2, random_gray_scale: float = 0.2, gaussian_blur: float = 0.0, - kernel_size: Optional[float] = None, sigmas: Tuple[float, float] = (0.1, 2), vf_prob: float = 0.0, hf_prob: float = 0.5, @@ -96,7 +91,6 @@ def __init__( min_scale=min_scale, random_gray_scale=random_gray_scale, gaussian_blur=gaussian_blur, - kernel_size=kernel_size, sigmas=sigmas, vf_prob=vf_prob, hf_prob=hf_prob, @@ -151,12 +145,8 @@ class MoCoV2Transform(SimCLRTransform): Probability of conversion to grayscale. gaussian_blur: Probability of Gaussian blur. - kernel_size: - Will be deprecated in favor of `sigmas` argument. If set, the old behavior applies and `sigmas` is ignored. - Used to calculate sigma of gaussian blur with kernel_size * input_size. sigmas: Tuple of min and max value from which the std of the gaussian kernel is sampled. - Is ignored if `kernel_size` is set. vf_prob: Probability that vertical flip is applied. hf_prob: @@ -186,7 +176,6 @@ def __init__( min_scale: float = 0.2, random_gray_scale: float = 0.2, gaussian_blur: float = 0.5, - kernel_size: Optional[float] = None, sigmas: Tuple[float, float] = (0.1, 2), vf_prob: float = 0.0, hf_prob: float = 0.5, @@ -205,7 +194,6 @@ def __init__( min_scale=min_scale, random_gray_scale=random_gray_scale, gaussian_blur=gaussian_blur, - kernel_size=kernel_size, sigmas=sigmas, vf_prob=vf_prob, hf_prob=hf_prob, diff --git a/lightly/transforms/msn_transform.py b/lightly/transforms/msn_transform.py index d49fe1b7f..8bf4ce5ab 100644 --- a/lightly/transforms/msn_transform.py +++ b/lightly/transforms/msn_transform.py @@ -1,4 +1,4 @@ -from typing import Dict, List, Optional, Tuple, Union +from typing import Dict, List, Tuple, Union from PIL.Image import Image from torch import Tensor @@ -60,12 +60,8 @@ class MSNTransform(MultiViewTransform): How much to jitter hue. gaussian_blur: Probability of Gaussian blur. - kernel_size: - Will be deprecated in favor of `sigmas` argument. If set, the old behavior applies and `sigmas` is ignored. - Used to calculate sigma of gaussian blur with kernel_size * input_size. sigmas: Tuple of min and max value from which the std of the gaussian kernel is sampled. - Is ignored if `kernel_size` is set. random_gray_scale: Probability of conversion to grayscale. hf_prob: @@ -91,7 +87,6 @@ def __init__( cj_sat: float = 0.8, cj_hue: float = 0.2, gaussian_blur: float = 0.5, - kernel_size: Optional[float] = None, sigmas: Tuple[float, float] = (0.1, 2), random_gray_scale: float = 0.2, hf_prob: float = 0.5, @@ -108,7 +103,6 @@ def __init__( cj_sat=cj_sat, cj_hue=cj_hue, gaussian_blur=gaussian_blur, - kernel_size=kernel_size, sigmas=sigmas, random_gray_scale=random_gray_scale, hf_prob=hf_prob, @@ -121,7 +115,6 @@ def __init__( cj_prob=cj_prob, cj_strength=cj_strength, gaussian_blur=gaussian_blur, - kernel_size=kernel_size, sigmas=sigmas, random_gray_scale=random_gray_scale, hf_prob=hf_prob, @@ -167,7 +160,6 @@ def __init__( cj_sat: float = 0.8, cj_hue: float = 0.2, gaussian_blur: float = 0.5, - kernel_size: Optional[float] = None, sigmas: Tuple[float, float] = (0.1, 2), random_gray_scale: float = 0.2, hf_prob: float = 0.5, @@ -187,11 +179,8 @@ def __init__( cj_sat: How much to jitter saturation. cj_hue: How much to jitter hue. gaussian_blur: Probability of Gaussian blur. - kernel_size: Will be deprecated in favor of `sigmas` argument. If set, - the old behavior applies and `sigmas` is ignored. Used to calculate - sigma of gaussian blur with kernel_size * input_size. sigmas: Tuple of min and max value from which the std of the gaussian - kernel is sampled. Is ignored if `kernel_size` is set. + kernel is sampled. random_gray_scale: Probability of conversion to grayscale. hf_prob: Probability that horizontal flip is applied. vf_prob: Probability that vertical flip is applied. @@ -211,7 +200,7 @@ def __init__( T.RandomVerticalFlip(p=vf_prob), T.RandomApply([color_jitter], p=cj_prob), T.RandomGrayscale(p=random_gray_scale), - GaussianBlur(kernel_size=kernel_size, sigmas=sigmas, prob=gaussian_blur), + GaussianBlur(sigmas=sigmas, prob=gaussian_blur), T.ToTensor(), T.Normalize(mean=normalize["mean"], std=normalize["std"]), ] diff --git a/lightly/transforms/simclr_transform.py b/lightly/transforms/simclr_transform.py index ee844d705..8d8c56d95 100644 --- a/lightly/transforms/simclr_transform.py +++ b/lightly/transforms/simclr_transform.py @@ -46,7 +46,6 @@ def __init__( min_scale: float = 0.08, random_gray_scale: float = 0.2, gaussian_blur: float = 0.5, - kernel_size: Optional[float] = None, sigmas: Tuple[float, float] = (0.1, 2), vf_prob: float = 0.0, hf_prob: float = 0.5, @@ -70,11 +69,8 @@ def __init__( min_scale: Minimum size of the randomized crop relative to the input_size. random_gray_scale: Probability of conversion to grayscale. gaussian_blur: Probability of Gaussian blur. - kernel_size: Will be deprecated in favor of `sigmas` argument. If set, - the old behavior applies and `sigmas` is ignored. Used to calculate - sigma of gaussian blur with kernel_size * input_size. sigmas: Tuple of min and max value from which the std of the gaussian - kernel is sampled. Is ignored if `kernel_size` is set. + kernel is sampled. vf_prob: Probability that vertical flip is applied. hf_prob: Probability that horizontal flip is applied. rr_prob: Probability that random rotation is applied. @@ -99,7 +95,6 @@ def __init__( min_scale=min_scale, random_gray_scale=random_gray_scale, gaussian_blur=gaussian_blur, - kernel_size=kernel_size, sigmas=sigmas, vf_prob=vf_prob, hf_prob=hf_prob, @@ -145,7 +140,6 @@ def __init__( min_scale: float = 0.08, random_gray_scale: float = 0.2, gaussian_blur: float = 0.5, - kernel_size: Optional[float] = None, sigmas: Tuple[float, float] = (0.1, 2), vf_prob: float = 0.0, hf_prob: float = 0.5, @@ -169,11 +163,8 @@ def __init__( min_scale: Minimum size of the randomized crop relative to the input_size. random_gray_scale: Probability of conversion to grayscale. gaussian_blur: Probability of Gaussian blur. - kernel_size: Will be deprecated in favor of `sigmas` argument. If set, - the old behavior applies and `sigmas` is ignored. Used to calculate - sigma of gaussian blur with kernel_size * input_size. sigmas: Tuple of min and max value from which the std of the gaussian - kernel is sampled. Is ignored if `kernel_size` is set. + kernel is sampled. vf_prob: Probability that vertical flip is applied. hf_prob: Probability that horizontal flip is applied. rr_prob: Probability that random rotation is applied. @@ -201,7 +192,7 @@ def __init__( random_rotation_transform(rr_prob=rr_prob, rr_degrees=rr_degrees), T.RandomApply([color_jitter], p=cj_prob), T.RandomGrayscale(p=random_gray_scale), - GaussianBlur(kernel_size=kernel_size, sigmas=sigmas, prob=gaussian_blur), + GaussianBlur(sigmas=sigmas, prob=gaussian_blur), T.ToTensor(), ] if normalize: diff --git a/lightly/transforms/simsiam_transform.py b/lightly/transforms/simsiam_transform.py index 0903638ef..800e71337 100644 --- a/lightly/transforms/simsiam_transform.py +++ b/lightly/transforms/simsiam_transform.py @@ -50,12 +50,8 @@ class SimSiamTransform(MultiViewTransform): Probability of conversion to grayscale. gaussian_blur: Probability of Gaussian blur. - kernel_size: - Will be deprecated in favor of `sigmas` argument. If set, the old behavior applies and `sigmas` is ignored. - Used to calculate sigma of gaussian blur with kernel_size * input_size. sigmas: Tuple of min and max value from which the std of the gaussian kernel is sampled. - Is ignored if `kernel_size` is set. vf_prob: Probability that vertical flip is applied. hf_prob: @@ -85,7 +81,6 @@ def __init__( min_scale: float = 0.2, random_gray_scale: float = 0.2, gaussian_blur: float = 0.5, - kernel_size: Optional[float] = None, sigmas: Tuple[float, float] = (0.1, 2), vf_prob: float = 0.0, hf_prob: float = 0.5, @@ -104,7 +99,6 @@ def __init__( min_scale=min_scale, random_gray_scale=random_gray_scale, gaussian_blur=gaussian_blur, - kernel_size=kernel_size, sigmas=sigmas, vf_prob=vf_prob, hf_prob=hf_prob, @@ -149,7 +143,6 @@ def __init__( min_scale: float = 0.2, random_gray_scale: float = 0.2, gaussian_blur: float = 0.5, - kernel_size: Optional[float] = None, sigmas: Tuple[float, float] = (0.1, 2), vf_prob: float = 0.0, hf_prob: float = 0.5, @@ -173,11 +166,8 @@ def __init__( min_scale: Minimum size of the randomized crop relative to the input_size. random_gray_scale: Probability of conversion to grayscale. gaussian_blur: Probability of Gaussian blur. - kernel_size: Will be deprecated in favor of `sigmas` argument. If set, - the old behavior applies and `sigmas` is ignored. Used to calculate - sigma of gaussian blur with kernel_size * input_size. sigmas: Tuple of min and max value from which the std of the gaussian - kernel is sampled. Is ignored if `kernel_size` is set. + kernel is sampled. vf_prob: Probability that vertical flip is applied. hf_prob: Probability that horizontal flip is applied. rr_prob: Probability that random rotation is applied. @@ -205,7 +195,7 @@ def __init__( random_rotation_transform(rr_prob=rr_prob, rr_degrees=rr_degrees), T.RandomApply([color_jitter], p=cj_prob), T.RandomGrayscale(p=random_gray_scale), - GaussianBlur(kernel_size=kernel_size, sigmas=sigmas, prob=gaussian_blur), + GaussianBlur(sigmas=sigmas, prob=gaussian_blur), T.ToTensor(), ] if normalize: diff --git a/lightly/transforms/smog_transform.py b/lightly/transforms/smog_transform.py index 288187b37..11960d987 100644 --- a/lightly/transforms/smog_transform.py +++ b/lightly/transforms/smog_transform.py @@ -1,4 +1,4 @@ -from typing import Dict, List, Optional, Tuple, Union +from typing import Dict, List, Tuple, Union from PIL.Image import Image from torch import Tensor @@ -39,8 +39,6 @@ class SMoGTransform(MultiViewTransform): Max_scales for each crop category. gaussian_blur_probs: Probability of Gaussian blur for each crop category. - gaussian_blur_kernel_sizes: - Deprecated values in favour of sigmas. gaussian_blur_sigmas: Tuple of min and max value from which the std of the gaussian kernel is sampled. solarize_probs: @@ -74,10 +72,6 @@ def __init__( crop_min_scales: Tuple[float, float] = (0.2, 0.05), crop_max_scales: Tuple[float, float] = (1.0, 0.2), gaussian_blur_probs: Tuple[float, float] = (0.5, 0.1), - gaussian_blur_kernel_sizes: Tuple[Optional[float], Optional[float]] = ( - None, - None, - ), gaussian_blur_sigmas: Tuple[float, float] = (0.1, 2), solarize_probs: Tuple[float, float] = (0.0, 0.2), hf_prob: float = 0.5, @@ -99,7 +93,6 @@ def __init__( crop_min_scale=crop_min_scales[i], crop_max_scale=crop_max_scales[i], gaussian_blur_prob=gaussian_blur_probs[i], - kernel_size=gaussian_blur_kernel_sizes[i], sigmas=gaussian_blur_sigmas, solarize_prob=solarize_probs[i], hf_prob=hf_prob, @@ -146,7 +139,6 @@ def __init__( crop_min_scale: float = 0.2, crop_max_scale: float = 1.0, gaussian_blur_prob: float = 0.5, - kernel_size: Optional[float] = None, sigmas: Tuple[float, float] = (0.1, 2), solarize_prob: float = 0.0, hf_prob: float = 0.5, @@ -166,11 +158,8 @@ def __init__( crop_min_scale: Minimum size of the randomized crop relative to the input_size. crop_max_scale: Maximum size of the randomized crop relative to the input_size. gaussian_blur_prob: Probability of Gaussian blur. - kernel_size: Will be deprecated in favor of `sigmas` argument. If set, - the old behavior applies and `sigmas` is ignored. Used to calculate - sigma of gaussian blur with kernel_size * input_size. sigmas: Tuple of min and max value from which the std of the gaussian - kernel is sampled. Is ignored if `kernel_size` is set. + kernel is sampled. solarize_prob: Probability of solarization. hf_prob: Probability that horizontal flip is applied. cj_prob: Probability that color jitter is applied. @@ -198,7 +187,6 @@ def __init__( T.RandomApply([color_jitter], p=cj_prob), T.RandomGrayscale(p=random_gray_scale), GaussianBlur( - kernel_size=kernel_size, prob=gaussian_blur_prob, sigmas=sigmas, ), diff --git a/lightly/transforms/swav_transform.py b/lightly/transforms/swav_transform.py index c9e8145e9..65e3ee91b 100644 --- a/lightly/transforms/swav_transform.py +++ b/lightly/transforms/swav_transform.py @@ -65,12 +65,8 @@ class SwaVTransform(MultiCropTranform): Probability of conversion to grayscale. gaussian_blur: Probability of Gaussian blur. - kernel_size: - Will be deprecated in favor of `sigmas` argument. If set, the old behavior applies and `sigmas` is ignored. - Used to calculate sigma of gaussian blur with kernel_size * input_size. sigmas: Tuple of min and max value from which the std of the gaussian kernel is sampled. - Is ignored if `kernel_size` is set. normalize: Dictionary with 'mean' and 'std' for torchvision.transforms.Normalize. @@ -94,7 +90,6 @@ def __init__( cj_hue: float = 0.2, random_gray_scale: float = 0.2, gaussian_blur: float = 0.5, - kernel_size: Optional[float] = None, sigmas: Tuple[float, float] = (0.1, 2), normalize: Union[None, Dict[str, List[float]]] = IMAGENET_NORMALIZE, ): @@ -111,7 +106,6 @@ def __init__( cj_hue=cj_hue, random_gray_scale=random_gray_scale, gaussian_blur=gaussian_blur, - kernel_size=kernel_size, sigmas=sigmas, normalize=normalize, ) @@ -160,7 +154,6 @@ def __init__( cj_hue: float = 0.2, random_gray_scale: float = 0.2, gaussian_blur: float = 0.5, - kernel_size: Optional[float] = None, sigmas: Tuple[float, float] = (0.1, 2), normalize: Union[None, Dict[str, List[float]]] = IMAGENET_NORMALIZE, ): @@ -185,11 +178,8 @@ def __init__( cj_hue: How much to jitter hue. random_gray_scale: Probability of conversion to grayscale. gaussian_blur: Probability of Gaussian blur. - kernel_size: Will be deprecated in favor of `sigmas` argument. If set, - the old behavior applies and `sigmas` is ignored. Used to calculate - sigma of gaussian blur with kernel_size * input_size. sigmas: Tuple of min and max value from which the std of the gaussian - kernel is sampled. Is ignored if `kernel_size` is set. + kernel is sampled. normalize: Dictionary with 'mean' and 'std' for torchvision.transforms.Normalize. @@ -208,7 +198,7 @@ def __init__( T.ColorJitter(), T.RandomApply([color_jitter], p=cj_prob), T.RandomGrayscale(p=random_gray_scale), - GaussianBlur(kernel_size=kernel_size, sigmas=sigmas, prob=gaussian_blur), + GaussianBlur(sigmas=sigmas, prob=gaussian_blur), T.ToTensor(), ] if normalize: diff --git a/lightly/transforms/vicreg_transform.py b/lightly/transforms/vicreg_transform.py index b794a7bd4..3a04b3120 100644 --- a/lightly/transforms/vicreg_transform.py +++ b/lightly/transforms/vicreg_transform.py @@ -55,12 +55,8 @@ class VICRegTransform(MultiViewTransform): Probability of solarization. gaussian_blur: Probability of Gaussian blur. - kernel_size: - Will be deprecated in favor of `sigmas` argument. If set, the old behavior applies and `sigmas` is ignored. - Used to calculate sigma of gaussian blur with kernel_size * input_size. sigmas: Tuple of min and max value from which the std of the gaussian kernel is sampled. - Is ignored if `kernel_size` is set. vf_prob: Probability that vertical flip is applied. hf_prob: @@ -91,7 +87,6 @@ def __init__( random_gray_scale: float = 0.2, solarize_prob: float = 0.1, gaussian_blur: float = 0.5, - kernel_size: Optional[float] = None, sigmas: Tuple[float, float] = (0.1, 2), vf_prob: float = 0.0, hf_prob: float = 0.5, @@ -111,7 +106,6 @@ def __init__( random_gray_scale=random_gray_scale, solarize_prob=solarize_prob, gaussian_blur=gaussian_blur, - kernel_size=kernel_size, sigmas=sigmas, vf_prob=vf_prob, hf_prob=hf_prob, @@ -158,7 +152,6 @@ def __init__( random_gray_scale: float = 0.2, solarize_prob: float = 0.1, gaussian_blur: float = 0.5, - kernel_size: Optional[float] = None, sigmas: Tuple[float, float] = (0.2, 2), vf_prob: float = 0.0, hf_prob: float = 0.5, @@ -181,11 +174,8 @@ def __init__( random_gray_scale: Probability of conversion to grayscale. solarize_prob: Probability of solarization. gaussian_blur: Probability of Gaussian blur. - kernel_size: Will be deprecated in favor of `sigmas` argument. If set, - the old behavior applies and `sigmas` is ignored. Used to calculate - sigma of gaussian blur with kernel_size * input_size. sigmas: Tuple of min and max value from which the std of the gaussian - kernel is sampled. Is ignored if `kernel_size` is set. + kernel is sampled. vf_prob: Probability that vertical flip is applied. hf_prob: Probability that horizontal flip is applied. rr_prob: Probability that random rotation is applied. @@ -214,7 +204,7 @@ def __init__( T.RandomApply([color_jitter], p=cj_prob), T.RandomGrayscale(p=random_gray_scale), RandomSolarization(prob=solarize_prob), - GaussianBlur(kernel_size=kernel_size, sigmas=sigmas, prob=gaussian_blur), + GaussianBlur(sigmas=sigmas, prob=gaussian_blur), T.ToTensor(), ] if normalize: diff --git a/lightly/transforms/vicregl_transform.py b/lightly/transforms/vicregl_transform.py index 2727a692d..0982e0403 100644 --- a/lightly/transforms/vicregl_transform.py +++ b/lightly/transforms/vicregl_transform.py @@ -1,4 +1,4 @@ -from typing import Dict, List, Optional, Tuple, Union +from typing import Dict, List, Tuple, Union from PIL.Image import Image from torch import Tensor @@ -53,23 +53,13 @@ class VICRegLTransform(ImageGridTransform): Probability of Gaussian blur for the global crop views. local_gaussian_blur_prob: Probability of Gaussian blur for the local crop views. - global_gaussian_blur_kernel_size: - Will be deprecated in favor of `global_gaussian_blur_sigmas` argument. - If set, the old behavior applies and `global_gaussian_blur_sigmas` - is ignored. Used to calculate sigma of gaussian blur with - global_gaussian_blur_kernel_size * input_size. Applied to global crop views. - local_gaussian_blur_kernel_size: - Will be deprecated in favor of `local_gaussian_blur_sigmas` argument. - If set, the old behavior applies and `local_gaussian_blur_sigmas` - is ignored. Used to calculate sigma of gaussian blur with - local_gaussian_blur_kernel_size * input_size. Applied to local crop views. global_gaussian_blur_sigmas: Tuple of min and max value from which the std of the gaussian kernel - is sampled. It is ignored if `global_gaussian_blur_kernel_size` is set. + is sampled. Applied to global crop views. local_gaussian_blur_sigmas: Tuple of min and max value from which the std of the gaussian kernel - is sampled. It is ignored if `local_gaussian_blur_kernel_size` is set. + is sampled. Applied to local crop views. global_solarize_prob: Probability of solarization for the global crop views. @@ -108,8 +98,6 @@ def __init__( local_grid_size: int = 3, global_gaussian_blur_prob: float = 0.5, local_gaussian_blur_prob: float = 0.1, - global_gaussian_blur_kernel_size: Optional[float] = None, - local_gaussian_blur_kernel_size: Optional[float] = None, global_gaussian_blur_sigmas: Tuple[float, float] = (0.1, 2), local_gaussian_blur_sigmas: Tuple[float, float] = (0.1, 2), global_solarize_prob: float = 0.0, @@ -136,7 +124,6 @@ def __init__( ), VICRegLViewTransform( gaussian_blur_prob=global_gaussian_blur_prob, - gaussian_blur_kernel_size=global_gaussian_blur_kernel_size, gaussian_blur_sigmas=global_gaussian_blur_sigmas, solarize_prob=global_solarize_prob, cj_prob=cj_prob, @@ -159,7 +146,6 @@ def __init__( ), VICRegLViewTransform( gaussian_blur_prob=local_gaussian_blur_prob, - gaussian_blur_kernel_size=local_gaussian_blur_kernel_size, gaussian_blur_sigmas=local_gaussian_blur_sigmas, solarize_prob=local_solarize_prob, cj_prob=cj_prob, @@ -201,7 +187,6 @@ class VICRegLViewTransform: def __init__( self, gaussian_blur_prob: float = 0.5, - gaussian_blur_kernel_size: Optional[float] = None, gaussian_blur_sigmas: Tuple[float, float] = (0.1, 2), solarize_prob: float = 0.0, cj_prob: float = 1.0, @@ -217,13 +202,8 @@ def __init__( Args: gaussian_blur_prob: Probability of Gaussian blur. - gaussian_blur_kernel_size: Will be deprecated in favor of - `gaussian_blur_sigmas` argument. If set, the old behavior applies - and `gaussian_blur_sigmas` is ignored. Used to calculate sigma of - gaussian blur with gaussian_blur_kernel_size * input_size. gaussian_blur_sigmas: Tuple of min and max value from which the std of - the gaussian kernel is sampled. Is ignored if - `gaussian_blur_kernel_size` is set. + the gaussian kernel is sampled. solarize_prob: Probability of solarization. cj_prob: Probability that color jitter is applied. cj_strength: Strength of the color jitter. `cj_bright`, `cj_contrast`, @@ -248,7 +228,6 @@ def __init__( T.RandomApply([color_jitter], p=cj_prob), T.RandomGrayscale(p=random_gray_scale), GaussianBlur( - kernel_size=gaussian_blur_kernel_size, prob=gaussian_blur_prob, sigmas=gaussian_blur_sigmas, ), diff --git a/lightly/transforms/wmse_transform.py b/lightly/transforms/wmse_transform.py index dc259bcc2..7db2d3f04 100644 --- a/lightly/transforms/wmse_transform.py +++ b/lightly/transforms/wmse_transform.py @@ -1,4 +1,4 @@ -from typing import Dict, List, Optional, Tuple +from typing import Dict, List, Tuple from lightly.transforms.gaussian_blur import GaussianBlur from lightly.transforms.multi_view_transform import MultiViewTransform @@ -54,12 +54,8 @@ class WMSETransform(MultiViewTransform): Probability that horizontal flip is applied. gaussian_blur: Probability of Gaussian blur. - kernel_size: - Will be deprecated in favor of `sigmas` argument. If set, the old behavior applies and `sigmas` is ignored. - Used to calculate sigma of gaussian blur with kernel_size * input_size. sigmas: Tuple of min and max value from which the std of the gaussian kernel is sampled. - Is ignored if `kernel_size` is set. normalize: Dictionary with 'mean' and 'std' for torchvision.transforms.Normalize. """ @@ -77,7 +73,6 @@ def __init__( random_gray_scale: float = 0.1, hf_prob: float = 0.5, gaussian_blur: float = 0.5, - kernel_size: Optional[int] = None, sigmas: Tuple[float, float] = (0.1, 2.0), normalize: Dict[str, List[float]] = IMAGENET_NORMALIZE, ): @@ -95,9 +90,7 @@ def __init__( interpolation=T.InterpolationMode.BICUBIC, ), T.RandomHorizontalFlip(p=hf_prob), - GaussianBlur( - kernel_size=kernel_size, sigmas=sigmas, prob=gaussian_blur - ), + GaussianBlur(sigmas=sigmas, prob=gaussian_blur), T.ToTensor(), T.Normalize(mean=normalize["mean"], std=normalize["std"]), ] diff --git a/tests/transforms/test_gaussian_blur.py b/tests/transforms/test_gaussian_blur.py index 171383649..88f329a16 100644 --- a/tests/transforms/test_gaussian_blur.py +++ b/tests/transforms/test_gaussian_blur.py @@ -20,10 +20,6 @@ def test_on_tensor(self, w: int, h: int) -> None: sample_tensor = torch.randn(3, h, w) gaussian_blur(sample_tensor) - def test_raise_kernel_size_deprecation(self) -> None: - with pytest.warns(DeprecationWarning): - GaussianBlur(kernel_size=2) - - def test_raise_scale_deprecation(self) -> None: - with pytest.warns(DeprecationWarning): - GaussianBlur(scale=0.1) + def test_init__keyword_only(self) -> None: + with pytest.raises(TypeError): + GaussianBlur(0.5) # type: ignore[misc] From 26f99dce13acfc9bb767475006e370610e8d24ae Mon Sep 17 00:00:00 2001 From: gabrielfruet Date: Tue, 11 Aug 2026 18:19:42 -0300 Subject: [PATCH 4/9] Remove remaining deprecated arguments and losses - MSNLoss(me_max_weight=): use regularization_weight - SymNegCosineSimilarityLoss: use NegativeCosineSimilarity - MemoryBankModule bare positive int size: pass (num_features, dim). size=0 still disables the bank, so NTXentLoss(memory_bank_size=0) is unaffected. Inferring the feature dimension from the first batch broke distributed training, so it now raises instead of warning, and the bank is allocated in __init__. - trainer.weights_summary: use trainer.enable_model_summary and summary_callback.max_depth Co-Authored-By: Claude Opus 5 --- docs/source/lightly.loss.rst | 3 - lightly/cli/config/config.yaml | 2 - lightly/cli/train_cli.py | 2 +- lightly/embedding/_base.py | 14 +--- lightly/embedding/callbacks.py | 51 +------------- lightly/loss/__init__.py | 1 - lightly/loss/directclr_loss.py | 4 -- lightly/loss/msn_loss.py | 22 ------ lightly/loss/ntx_ent_loss.py | 4 -- lightly/loss/regularizer/co2.py | 5 +- lightly/loss/sym_neg_cos_sim_loss.py | 87 ------------------------ lightly/models/modules/memory_bank.py | 33 +++------ lightly/models/modules/nn_memory_bank.py | 8 +-- tests/cli/test_cli_get_lighty_config.py | 2 +- tests/embedding/test_callbacks.py | 36 ---------- tests/loss/test_msn_loss.py | 4 -- tests/loss/test_ntx_ent_loss.py | 6 +- tests/loss/test_sym_neg_cos_sim_loss.py | 56 --------------- tests/models/modules/test_memory_bank.py | 42 +++++------- 19 files changed, 38 insertions(+), 344 deletions(-) delete mode 100644 lightly/loss/sym_neg_cos_sim_loss.py delete mode 100644 tests/loss/test_sym_neg_cos_sim_loss.py diff --git a/docs/source/lightly.loss.rst b/docs/source/lightly.loss.rst index 437d6e077..734437abe 100644 --- a/docs/source/lightly.loss.rst +++ b/docs/source/lightly.loss.rst @@ -77,9 +77,6 @@ lightly.loss .. autoclass:: lightly.loss.swav_loss.SwaVLoss :members: -.. autoclass:: lightly.loss.sym_neg_cos_sim_loss.SymNegCosineSimilarityLoss - :members: - .. autoclass:: lightly.loss.tico_loss.TiCoLoss :members: diff --git a/lightly/cli/config/config.yaml b/lightly/cli/config/config.yaml index 40e879ada..2484adf92 100644 --- a/lightly/cli/config/config.yaml +++ b/lightly/cli/config/config.yaml @@ -68,8 +68,6 @@ trainer: max_epochs: 100 # Number of epochs to train for. precision: 32 # If set to 16, will use half-precision. enable_model_summary: True # Whether to enable model summarisation. - weights_summary: # [deprecated] Use enable_model_summary - # and summary_callback.max_depth. # checkpoint_callback namespace: Modify the checkpoint callback checkpoint_callback: diff --git a/lightly/cli/train_cli.py b/lightly/cli/train_cli.py index 9a343a6ee..ee2c898bf 100644 --- a/lightly/cli/train_cli.py +++ b/lightly/cli/train_cli.py @@ -177,7 +177,7 @@ def train_cli(cfg): >>> lightly-ssl-train input_dir=data/ trainer.max_epochs=10 >>> >>> # print a full summary of the model - >>> lightly-ssl-train input_dir=data/ trainer.weights_summary=full + >>> lightly-ssl-train input_dir=data/ summary_callback.max_depth=-1 """ return _train_cli(cfg) diff --git a/lightly/embedding/_base.py b/lightly/embedding/_base.py index b7e45c71b..b18fe9c15 100644 --- a/lightly/embedding/_base.py +++ b/lightly/embedding/_base.py @@ -2,11 +2,9 @@ # Copyright (c) 2020. Lightly AG and its affiliates. # All Rights Reserved -import copy import os from typing import Any, List, Optional, Sequence, Tuple, Union -import omegaconf from omegaconf import DictConfig from pytorch_lightning import LightningModule, Trainer from pytorch_lightning.callbacks import Callback @@ -92,7 +90,6 @@ def train_embedding( max_epochs: (int) Maximum number of epochs to train gpus: (int) Number of gpus to use enable_model_summary: (bool) Whether to enable model summarisation. - weights_summary: (str) DEPRECATED. How to print a summary of the model and weights. checkpoint_callback_config: ModelCheckpoint callback arguments summary_callback_config: ModelSummary callback arguments @@ -108,19 +105,12 @@ def train_embedding( trainer_callbacks.append(checkpoint_cb) summary_cb = callbacks.create_summary_callback( - summary_callback_config=summary_callback_config, - trainer_config=trainer_config, + summary_callback_config=summary_callback_config ) if summary_cb is not None: trainer_callbacks.append(summary_cb) - # Remove weights_summary from trainer_config now that the summary callback - # has been created. TODO: Drop support for the "weights_summary" argument. - trainer_config_copy = copy.deepcopy(trainer_config) - if "weights_summary" in trainer_config_copy: - with omegaconf.open_dict(trainer_config_copy): - del trainer_config_copy["weights_summary"] - trainer = Trainer(**trainer_config_copy, callbacks=trainer_callbacks) # type: ignore[misc] + trainer = Trainer(**trainer_config, callbacks=trainer_callbacks) # type: ignore[misc] trainer.fit(self) diff --git a/lightly/embedding/callbacks.py b/lightly/embedding/callbacks.py index e4d620af4..9ec4e5fd3 100644 --- a/lightly/embedding/callbacks.py +++ b/lightly/embedding/callbacks.py @@ -4,8 +4,6 @@ from omegaconf import DictConfig from pytorch_lightning.callbacks import ModelCheckpoint, ModelSummary -from lightly.utils.hipify import print_as_warning - def create_checkpoint_callback( save_last: bool = False, @@ -39,27 +37,18 @@ def create_checkpoint_callback( ) -def create_summary_callback( - summary_callback_config: DictConfig, trainer_config: DictConfig -) -> ModelSummary: +def create_summary_callback(summary_callback_config: DictConfig) -> ModelSummary: """Creates a model summary callback based on the configuration. Args: summary_callback_config: Configuration dictionary for the summary callback. - trainer_config: - Trainer configuration dictionary, which may include deprecated `weights_summary`. Returns: ModelSummary: The model summary callback. """ - # TODO: Drop support for the "weights_summary" argument. - weights_summary = trainer_config.get("weights_summary", None) - if weights_summary not in [None, "None"]: - return _create_summary_callback_deprecated(weights_summary) - else: - return _create_summary_callback(summary_callback_config["max_depth"]) + return _create_summary_callback(summary_callback_config["max_depth"]) def _create_summary_callback(max_depth: int) -> ModelSummary: @@ -76,39 +65,3 @@ def _create_summary_callback(max_depth: int) -> ModelSummary: """ return ModelSummary(max_depth=max_depth) - - -def _create_summary_callback_deprecated(weights_summary: str) -> ModelSummary: - """Constructs summary callback from the deprecated ``weights_summary`` argument. - - The ``weights_summary`` trainer argument was deprecated with the release - of pytorch lightning 1.7 in 08/2022. Support for this will be removed - in the future. - - Args: - weights_summary: - The deprecated `weights_summary` argument value ("top" or "full"). - - Returns: - ModelSummary: The initialized model summary callback based on the `weights_summary` argument. - - Raises: - ValueError: If an invalid value is provided for `weights_summary`. - - """ - print_as_warning( - "The configuration parameter 'trainer.weights_summary' is deprecated." - " Please use 'trainer.weights_summary: True' and set" - " 'checkpoint_callback.max_depth' to value 1 for the option 'top'" - " or -1 for the option 'full'." - ) - if weights_summary == "top": - max_depth = 1 - elif weights_summary == "full": - max_depth = -1 - else: - raise ValueError( - "Invalid value for the deprecated trainer.weights_summary" - " configuration parameter." - ) - return _create_summary_callback(max_depth=max_depth) diff --git a/lightly/loss/__init__.py b/lightly/loss/__init__.py index 453349d33..42142511a 100644 --- a/lightly/loss/__init__.py +++ b/lightly/loss/__init__.py @@ -22,7 +22,6 @@ from lightly.loss.patch_kernel_alignment_loss import PatchKernelAlignmentLoss from lightly.loss.pmsn_loss import PMSNCustomLoss, PMSNLoss from lightly.loss.swav_loss import SwaVLoss -from lightly.loss.sym_neg_cos_sim_loss import SymNegCosineSimilarityLoss from lightly.loss.tico_loss import TiCoLoss from lightly.loss.vicreg_loss import VICRegLoss from lightly.loss.vicregl_loss import VICRegLLoss diff --git a/lightly/loss/directclr_loss.py b/lightly/loss/directclr_loss.py index 9151db2b9..bc2fd2af9 100644 --- a/lightly/loss/directclr_loss.py +++ b/lightly/loss/directclr_loss.py @@ -30,10 +30,6 @@ class DirectCLRLoss(NTXentLoss): num_features are the number of negative samples stored in the memory bank. If num_features is 0, the memory bank is disabled. Use 0 for SimCLR. For MoCo we typically use numbers like 4096 or 65536. - Deprecated: If only a single integer is passed, it is interpreted as the - number of features and the feature dimension is inferred from the first - batch stored in the memory bank. Leaving out the feature dimension might - lead to errors in distributed training. gather_distributed: From NTXentLoss: if True then negatives from all GPUs are gathered before the loss calculation. If a memory bank is used and gather_distributed is diff --git a/lightly/loss/msn_loss.py b/lightly/loss/msn_loss.py index 55808f313..8142f1a5f 100644 --- a/lightly/loss/msn_loss.py +++ b/lightly/loss/msn_loss.py @@ -1,6 +1,4 @@ import math -import warnings -from typing import Optional import torch import torch.distributed as dist @@ -120,11 +118,6 @@ class MSNLoss(nn.Module): regularization_weight: Weight factor lambda by which the regularization loss is scaled. Set to 0 to disable regularization. - me_max_weight: - Deprecated, use `regularization_weight` instead. Takes precedence over - `regularization_weight` if not None. Weight factor lambda by which the mean - entropy maximization regularization loss is scaled. Set to 0 to disable - mean entropy maximization reguliarization. gather_distributed: If True, then target probabilities are gathered from all GPUs. @@ -149,7 +142,6 @@ def __init__( temperature: float = 0.1, sinkhorn_iterations: int = 3, regularization_weight: float = 1.0, - me_max_weight: Optional[float] = None, gather_distributed: bool = False, ): """Initializes the MSNLoss module with the specified parameters. @@ -161,11 +153,6 @@ def __init__( Number of sinkhorn normalization iterations on the targets. regularization_weight: Weight factor lambda by which the regularization loss is scaled. Set to 0 to disable regularization. - me_max_weight: - Deprecated, use `regularization_weight` instead. Takes precedence over - `regularization_weight` if not None. Weight factor lambda by which the mean - entropy maximization regularization loss is scaled. Set to 0 to disable mean - entropy maximization regularization. gather_distributed: If True, then target probabilities are gathered from all GPUs. @@ -191,15 +178,6 @@ def __init__( self.temperature = temperature self.sinkhorn_iterations = sinkhorn_iterations self.regularization_weight = regularization_weight - # Set regularization_weight to me_max_weight for backwards compatibility - if me_max_weight is not None: - warnings.warn( - DeprecationWarning( - "me_max_weight is deprecated in favor of regularization_weight and " - "will be removed in the future." - ) - ) - self.regularization_weight = me_max_weight self.gather_distributed = gather_distributed def forward( diff --git a/lightly/loss/ntx_ent_loss.py b/lightly/loss/ntx_ent_loss.py index 62f82621f..558c0515c 100644 --- a/lightly/loss/ntx_ent_loss.py +++ b/lightly/loss/ntx_ent_loss.py @@ -31,10 +31,6 @@ class NTXentLoss(nn.Module): number of negative samples stored in the memory bank. If num_features is 0, the memory bank is disabled. Use 0 for SimCLR. For MoCo we typically use numbers like 4096 or 65536. - Deprecated: If only a single integer is passed, it is interpreted as the - number of features and the feature dimension is inferred from the first - batch stored in the memory bank. Leaving out the feature dimension might - lead to errors in distributed training. gather_distributed: If True then negatives from all GPUs are gathered before the loss calculation. If a memory bank is used and gather_distributed is True, diff --git a/lightly/loss/regularizer/co2.py b/lightly/loss/regularizer/co2.py index d35587ada..5bbcca47c 100644 --- a/lightly/loss/regularizer/co2.py +++ b/lightly/loss/regularizer/co2.py @@ -25,10 +25,7 @@ class CO2Regularizer(Module): memory_bank_size: Size of the memory bank as (num_features, dim) tuple. num_features is the number of negatives stored in the bank. If set to 0, the memory bank is - disabled. Deprecated: If only a single integer is passed, it is interpreted - as the number of features and the feature dimension is inferred from the - first batch stored in the memory bank. Leaving out the feature dimension - might lead to errors in distributed training. + disabled. Examples: >>> # initialize loss function for MoCo diff --git a/lightly/loss/sym_neg_cos_sim_loss.py b/lightly/loss/sym_neg_cos_sim_loss.py deleted file mode 100644 index 87557388b..000000000 --- a/lightly/loss/sym_neg_cos_sim_loss.py +++ /dev/null @@ -1,87 +0,0 @@ -"""Symmetrized Negative Cosine Similarity Loss Functions""" - -# Copyright (c) 2020. Lightly AG and its affiliates. -# All Rights Reserved - -import warnings - -import torch -from torch import Tensor -from torch.nn import Module - - -class SymNegCosineSimilarityLoss(Module): - """Implementation of the Symmetrized Loss used in the SimSiam[0] paper. - - - [0] SimSiam, 2020, https://arxiv.org/abs/2011.10566 - - Examples: - >>> # initialize loss function - >>> loss_fn = SymNegCosineSimilarityLoss() - >>> - >>> # generate two random transforms of images - >>> t0 = transforms(images) - >>> t1 = transforms(images) - >>> - >>> # feed through SimSiam model - >>> out0, out1 = model(t0, t1) - >>> - >>> # calculate loss - >>> loss = loss_fn(out0, out1) - """ - - def __init__(self) -> None: - """Initializes the SymNegCosineSimilarityLoss module. - - Note: - SymNegCosineSimilarityLoss will be deprecated in favor of NegativeCosineSimilarity in the future. - """ - super().__init__() - warnings.warn( - Warning( - "SymNegCosineSimiliarityLoss will be deprecated in favor of " - + "NegativeCosineSimilarity in the future." - ), - DeprecationWarning, - ) - - def forward(self, out0: Tensor, out1: Tensor) -> Tensor: - """Forward pass through Symmetric Loss. - - Args: - out0: - Output projections of the first set of transformed images. - Expects the tuple to be of the form (z0, p0), where z0 is - the output of the backbone and projection MLP, and p0 is the - output of the prediction head. - out1: - Output projections of the second set of transformed images. - Expects the tuple to be of the form (z1, p1), where z1 is - the output of the backbone and projection MLP, and p1 is the - output of the prediction head. - - Returns: - Negative Cosine Similarity loss value. - """ - z0, p0 = out0 - z1, p1 = out1 - - loss: Tensor = ( - self._neg_cosine_simililarity(p0, z1) / 2 - + self._neg_cosine_simililarity(p1, z0) / 2 - ) - - return loss - - def _neg_cosine_simililarity(self, x: Tensor, y: Tensor) -> Tensor: - """Calculates the negative cosine similarity between two tensors. - - Args: - x: First input tensor. - y: Second input tensor. - - Returns: - Negative cosine similarity value. - """ - v = -torch.nn.functional.cosine_similarity(x, y.detach(), dim=-1).mean() - return v diff --git a/lightly/models/modules/memory_bank.py b/lightly/models/modules/memory_bank.py index 7c7606abd..ae9b1eda4 100644 --- a/lightly/models/modules/memory_bank.py +++ b/lightly/models/modules/memory_bank.py @@ -3,7 +3,6 @@ # Copyright (c) 2020. Lightly AG and its affiliates. # All Rights Reserved -import warnings from typing import Optional, Sequence, Tuple, Union import torch @@ -24,11 +23,8 @@ class MemoryBankModule(Module): Attributes: size: Size of the memory bank as (num_features, dim) tuple. If num_features is 0 - then the memory bank is disabled. Deprecated: If only a single integer is - passed, it is interpreted as the number of features and the feature - dimension is inferred from the first batch stored in the memory bank. - Leaving out the feature dimension might lead to errors in distributed - training. + then the memory bank is disabled. Pass 0 to disable the memory bank; any + other bare integer is rejected because the feature dimension is required. gather_distributed: If True then negatives from all gpus are gathered before the memory bank is updated. This results in more frequent updates of the memory bank and @@ -57,7 +53,7 @@ class MemoryBankModule(Module): def __init__( self, - size: Union[int, Sequence[int]] = 65536, + size: Union[int, Sequence[int]] = (65536, 128), gather_distributed: bool = False, feature_dim_first: bool = True, ): @@ -68,6 +64,12 @@ def __init__( raise ValueError( f"Illegal memory bank size {size}, all entries must be non-negative." ) + if isinstance(size, int) and size > 0: + raise ValueError( + f"Memory bank size 'size={size}' does not specify the feature " + f"dimension. Set it with 'size=({size}, dim)', or pass 'size=0' to " + "disable the memory bank." + ) self.size = size_tuple self.gather_distributed = gather_distributed @@ -85,17 +87,7 @@ def __init__( persistent=False, ) - if isinstance(size, int) and size > 0: - warnings.warn( - ( - f"Memory bank size 'size={size}' does not specify feature " - "dimension. It is recommended to set the feature dimension with " - "'size=(n, dim)' when creating the memory bank. Distributed " - "training might fail if the feature dimension is not set." - ), - UserWarning, - ) - elif len(size_tuple) > 1: + if len(size_tuple) > 1: self._init_memory_bank(size=size_tuple) def forward( @@ -125,11 +117,6 @@ def forward( if self.size[0] == 0: return output, None - # Initialize the memory bank if it is not already done. - if self.bank.ndim == 1: - dim = output.shape[1:] - self._init_memory_bank(size=(*self.size, *dim)) - # query and update memory bank bank = self.bank.clone().detach() if self.feature_dim_first: diff --git a/lightly/models/modules/nn_memory_bank.py b/lightly/models/modules/nn_memory_bank.py index 0392bd24d..a264ce161 100644 --- a/lightly/models/modules/nn_memory_bank.py +++ b/lightly/models/modules/nn_memory_bank.py @@ -23,11 +23,7 @@ class NNMemoryBankModule(MemoryBankModule): Attributes: size: Size of the memory bank as (num_features, dim) tuple. If num_features is 0 - then the memory bank is disabled. Deprecated: If only a single integer is - passed, it is interpreted as the number of features and the feature - dimension is inferred from the first batch stored in the memory bank. - Leaving out the feature dimension might lead to errors in distributed - training. + then the memory bank is disabled. Examples: >>> model = NNCLR(backbone) @@ -44,7 +40,7 @@ class NNMemoryBankModule(MemoryBankModule): """ - def __init__(self, size: Union[int, Sequence[int]] = 2**16): + def __init__(self, size: Union[int, Sequence[int]] = (2**16, 128)): super(NNMemoryBankModule, self).__init__(size) def forward( # type: ignore[override] # TODO(Philipp, 11/23): Fix signature to match parent class. diff --git a/tests/cli/test_cli_get_lighty_config.py b/tests/cli/test_cli_get_lighty_config.py index e7aa38370..25b92e8a8 100644 --- a/tests/cli/test_cli_get_lighty_config.py +++ b/tests/cli/test_cli_get_lighty_config.py @@ -6,5 +6,5 @@ def test_get_lightly_config() -> None: # Assert some default values assert conf.checkpoint == "" assert conf.loader.batch_size == 16 - assert conf.trainer.weights_summary is None + assert conf.trainer.enable_model_summary is True assert conf.summary_callback.max_depth == 1 diff --git a/tests/embedding/test_callbacks.py b/tests/embedding/test_callbacks.py index dc57bbabe..87f0fa8cd 100644 --- a/tests/embedding/test_callbacks.py +++ b/tests/embedding/test_callbacks.py @@ -1,4 +1,3 @@ -import pytest from omegaconf import OmegaConf from lightly.embedding import callbacks @@ -7,40 +6,5 @@ def test_create_summary_callback(): summary_cb = callbacks.create_summary_callback( summary_callback_config=OmegaConf.create({"max_depth": 99}), - trainer_config=OmegaConf.create(), ) assert summary_cb._max_depth == 99 - - -def test_create_summary_callback__weights_summary(): - # If "weights_summary" is specified, it takes precedence. - summary_cb = callbacks.create_summary_callback( - summary_callback_config=OmegaConf.create({"max_depth": 99}), - trainer_config=OmegaConf.create({"weights_summary": "top"}), - ) - assert summary_cb._max_depth == 1 - - summary_cb = callbacks.create_summary_callback( - summary_callback_config=OmegaConf.create({"max_depth": 99}), - trainer_config=OmegaConf.create({"weights_summary": "full"}), - ) - assert summary_cb._max_depth == -1 - - # If "weights_summary" is None or "None", normal config is applied. - summary_cb = callbacks.create_summary_callback( - summary_callback_config=OmegaConf.create({"max_depth": 99}), - trainer_config=OmegaConf.create({"weights_summary": None}), - ) - assert summary_cb._max_depth == 99 - - summary_cb = callbacks.create_summary_callback( - summary_callback_config=OmegaConf.create({"max_depth": 99}), - trainer_config=OmegaConf.create({"weights_summary": "None"}), - ) - assert summary_cb._max_depth == 99 - - with pytest.raises(ValueError): - callbacks.create_summary_callback( - summary_callback_config=OmegaConf.create(), - trainer_config=OmegaConf.create({"weights_summary": "invalid"}), - ) diff --git a/tests/loss/test_msn_loss.py b/tests/loss/test_msn_loss.py index 7574fed30..0a45a3d64 100644 --- a/tests/loss/test_msn_loss.py +++ b/tests/loss/test_msn_loss.py @@ -39,10 +39,6 @@ def test__init__sinkhorn_iterations(self) -> None: with pytest.raises(ValueError): MSNLoss(sinkhorn_iterations=-1) - def test__init__me_max_weight(self) -> None: - criterion = MSNLoss(regularization_weight=0.0, me_max_weight=0.5) - assert criterion.regularization_weight == 0.5 - def test_prototype_probabilitiy(self) -> None: torch.manual_seed(0) queries = F.normalize(torch.rand((8, 10)), dim=1) diff --git a/tests/loss/test_ntx_ent_loss.py b/tests/loss/test_ntx_ent_loss.py index c717097c0..ec46074fe 100644 --- a/tests/loss/test_ntx_ent_loss.py +++ b/tests/loss/test_ntx_ent_loss.py @@ -74,7 +74,7 @@ def test_with_correlated_embedding( out0.requires_grad = True loss_function = NTXentLoss( temperature=temperature, - memory_bank_size=memory_bank_size, + memory_bank_size=(memory_bank_size, 2) if memory_bank_size > 0 else 0, gather_distributed=gather_distributed, ) if memory_bank_size > 0: @@ -121,7 +121,7 @@ def test_forward_pass_neg_temp(self) -> None: assert (l1 - l2).pow(2).item() == pytest.approx(0.0) def test_forward_pass_memory_bank(self) -> None: - loss = NTXentLoss(memory_bank_size=64) + loss = NTXentLoss(memory_bank_size=(64, 32)) for bsz in range(1, 20): batch_1 = torch.randn((bsz, 32)) batch_2 = torch.randn((bsz, 32)) @@ -129,7 +129,7 @@ def test_forward_pass_memory_bank(self) -> None: @pytest.mark.skipif(not torch.cuda.is_available(), reason="No cuda") def test_forward_pass_memory_bank_cuda(self) -> None: - loss = NTXentLoss(memory_bank_size=64) + loss = NTXentLoss(memory_bank_size=(64, 32)) for bsz in range(1, 20): batch_1 = torch.randn((bsz, 32)).cuda() batch_2 = torch.randn((bsz, 32)).cuda() diff --git a/tests/loss/test_sym_neg_cos_sim_loss.py b/tests/loss/test_sym_neg_cos_sim_loss.py deleted file mode 100644 index 5b0d7d923..000000000 --- a/tests/loss/test_sym_neg_cos_sim_loss.py +++ /dev/null @@ -1,56 +0,0 @@ -import pytest -import torch - -from lightly.loss import SymNegCosineSimilarityLoss - - -class TestSymNegCosineSimilarityLoss: - @pytest.mark.parametrize("bsz", range(1, 20)) - def test_forward_pass(self, bsz: int) -> None: - loss = SymNegCosineSimilarityLoss() - z0 = torch.randn((bsz, 32)) - p0 = torch.randn((bsz, 32)) - z1 = torch.randn((bsz, 32)) - p1 = torch.randn((bsz, 32)) - - # symmetry - l1 = loss((z0, p0), (z1, p1)) - l2 = loss((z1, p1), (z0, p0)) - assert torch.allclose(l1, l2) - - @pytest.mark.skipif(not torch.cuda.is_available(), reason="No cuda") - @pytest.mark.parametrize("bsz", range(1, 20)) - def test_forward_pass_cuda(self, bsz: int) -> None: - loss = SymNegCosineSimilarityLoss() - z0 = torch.randn((bsz, 32)).cuda() - p0 = torch.randn((bsz, 32)).cuda() - z1 = torch.randn((bsz, 32)).cuda() - p1 = torch.randn((bsz, 32)).cuda() - - # symmetry - l1 = loss((z0, p0), (z1, p1)) - l2 = loss((z1, p1), (z0, p0)) - assert torch.allclose(l1, l2) - - @pytest.mark.parametrize("bsz", range(1, 20)) - def test_neg_cosine_simililarity(self, bsz: int) -> None: - loss = SymNegCosineSimilarityLoss() - x = torch.randn((bsz, 32)) - y = torch.randn((bsz, 32)) - - # symmetry - l1 = loss._neg_cosine_simililarity(x, y) - l2 = loss._neg_cosine_simililarity(y, x) - assert torch.allclose(l1, l2) - - @pytest.mark.skipif(not torch.cuda.is_available(), reason="No cuda") - @pytest.mark.parametrize("bsz", range(1, 20)) - def test_neg_cosine_simililarity_cuda(self, bsz: int) -> None: - loss = SymNegCosineSimilarityLoss() - x = torch.randn((bsz, 32)).cuda() - y = torch.randn((bsz, 32)).cuda() - - # symmetry - l1 = loss._neg_cosine_simililarity(x, y) - l2 = loss._neg_cosine_simililarity(y, x) - assert torch.allclose(l1, l2) diff --git a/tests/models/modules/test_memory_bank.py b/tests/models/modules/test_memory_bank.py index fb902f3e4..d7e0e3e4f 100644 --- a/tests/models/modules/test_memory_bank.py +++ b/tests/models/modules/test_memory_bank.py @@ -15,7 +15,7 @@ def test_forward_easy(self) -> None: bsz = 3 dim, size = 2, 9 n = 33 * bsz - memory_bank = MemoryBankModule(size=size) + memory_bank = MemoryBankModule(size=(size, dim)) ptr = 0 for i in range(0, n, bsz): @@ -40,7 +40,7 @@ def test_forward(self) -> None: bsz = 3 dim, size = 2, 10 n = 33 * bsz - memory_bank = MemoryBankModule(size=size) + memory_bank = MemoryBankModule(size=(size, dim)) for i in range(0, n, bsz): # see if there are any problems when the bank size @@ -53,7 +53,7 @@ def test_forward__cuda(self) -> None: bsz = 3 dim, size = 2, 10 n = 33 * bsz - memory_bank = MemoryBankModule(size=size) + memory_bank = MemoryBankModule(size=(size, dim)) device = torch.device("cuda") memory_bank.to(device=device) @@ -80,18 +80,24 @@ def test_init__negative_size(self) -> None: ): MemoryBankModule(size=(10, -1)) - def test_init__no_dim_warning(self) -> None: - with pytest.warns( - UserWarning, + def test_init__no_dim(self) -> None: + with pytest.raises( + ValueError, match=re.escape( - "Memory bank size 'size=10' does not specify feature " - "dimension. It is recommended to set the feature dimension with " - "'size=(n, dim)' when creating the memory bank. Distributed " - "training might fail if the feature dimension is not set." + "Memory bank size 'size=10' does not specify the feature dimension. " + "Set it with 'size=(10, dim)', or pass 'size=0' to disable the " + "memory bank." ), ): MemoryBankModule(size=10) + def test_init__disabled(self) -> None: + memory_bank = MemoryBankModule(size=0) + x = torch.randn(3, 2) + out, bank = memory_bank(x, update=True) + assert out.tolist() == x.tolist() + assert bank is None + def test_forward(self) -> None: torch.manual_seed(0) memory_bank = MemoryBankModule(size=(5, 2), feature_dim_first=False) @@ -132,22 +138,6 @@ def test_forward(self) -> None: # Verify that memory bank is overwritten. assert memory_bank.bank[:3].tolist() == x2.tolist() - def test_forward__no_dim(self) -> None: - torch.manual_seed(0) - # Only specify size but not feature dimension. - memory_bank = MemoryBankModule(size=5, feature_dim_first=False) - x0 = torch.randn(3, 2) - out0, bank0 = memory_bank(x0, update=True) - # Verify that output is same as input. - assert out0.tolist() == x0.tolist() - # Verify that memory bank was initialized and has correct shape. - assert bank0.shape == (5, 2) - assert memory_bank.bank.shape == (5, 2) - # Verify that output bank does not contain features from x0. - assert bank0[:3].tolist() != x0.tolist() - # Verify that memory bank was updated. - assert memory_bank.bank[:3].tolist() == x0.tolist() - def test_forward__dim_first(self) -> None: torch.manual_seed(0) memory_bank = MemoryBankModule(size=(5, 2), feature_dim_first=True) From 04481d4f81fdca6fedb9f80e1928d86b1b83938f Mon Sep 17 00:00:00 2001 From: gabrielfruet Date: Tue, 8 Sep 2026 11:03:39 -0300 Subject: [PATCH 5/9] fix: remove undefined cast() call from GaussianBlur The master merge kept the pre-refactor __call__ body while taking the new import line, so every Tensor input raised NameError. Co-Authored-By: Claude Opus 5 --- lightly/transforms/gaussian_blur.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/lightly/transforms/gaussian_blur.py b/lightly/transforms/gaussian_blur.py index fc7497789..993dc59e9 100644 --- a/lightly/transforms/gaussian_blur.py +++ b/lightly/transforms/gaussian_blur.py @@ -49,9 +49,10 @@ def __call__(self, sample: Union[Tensor, Image]) -> Union[Tensor, Image]: prob = np.random.random_sample() # Convert to PIL image if it's a tensor, otherwise use as is + sample_pil: Image if isinstance(sample, Tensor): is_input_tensor = True - sample_pil = cast(Image, F.to_pil_image(sample)) + sample_pil = F.to_pil_image(sample) else: is_input_tensor = False sample_pil = sample From 56b55892132859e75fc07fdd5596f091f270ca77 Mon Sep 17 00:00:00 2001 From: gabrielfruet Date: Tue, 8 Sep 2026 11:03:39 -0300 Subject: [PATCH 6/9] fix(docs): pass feature dim to memory banks in benchmark scripts MemoryBankModule now rejects a bare int size. Co-Authored-By: Claude Opus 5 --- docs/source/getting_started/benchmarks/cifar10_benchmark.py | 4 ++-- .../source/getting_started/benchmarks/imagenette_benchmark.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/docs/source/getting_started/benchmarks/cifar10_benchmark.py b/docs/source/getting_started/benchmarks/cifar10_benchmark.py index d97faafdc..a89f11462 100644 --- a/docs/source/getting_started/benchmarks/cifar10_benchmark.py +++ b/docs/source/getting_started/benchmarks/cifar10_benchmark.py @@ -310,7 +310,7 @@ def __init__(self, dataloader_kNN, num_classes): # create our loss with the optional memory bank self.criterion = NTXentLoss( temperature=0.1, - memory_bank_size=4096, + memory_bank_size=(4096, 128), ) def forward(self, x): @@ -662,7 +662,7 @@ def __init__(self, dataloader_kNN, num_classes): ) self.criterion = NTXentLoss() - self.memory_bank = modules.NNMemoryBankModule(size=4096) + self.memory_bank = modules.NNMemoryBankModule(size=(4096, 256)) def forward(self, x): y = self.backbone(x).flatten(start_dim=1) diff --git a/docs/source/getting_started/benchmarks/imagenette_benchmark.py b/docs/source/getting_started/benchmarks/imagenette_benchmark.py index 67299677f..b08453f3f 100644 --- a/docs/source/getting_started/benchmarks/imagenette_benchmark.py +++ b/docs/source/getting_started/benchmarks/imagenette_benchmark.py @@ -581,7 +581,7 @@ def __init__(self, dataloader_kNN, num_classes): self.prediction_head = heads.NNCLRPredictionHead(256, 4096, 256) self.criterion = NTXentLoss() - self.memory_bank = modules.NNMemoryBankModule(size=4096) + self.memory_bank = modules.NNMemoryBankModule(size=(4096, 256)) def forward(self, x): y = self.backbone(x).flatten(start_dim=1) From dcab8b3e895044a7a66e9b06b8d2d81a6c5eb19b Mon Sep 17 00:00:00 2001 From: gabrielfruet Date: Thu, 10 Sep 2026 18:30:49 -0300 Subject: [PATCH 7/9] fix: require explicit size for MemoryBankModule The (65536, 128) default silently built a 128-dim bank, so a mismatched feature dim died in einsum instead of raising the size error. Co-Authored-By: Claude Opus 5 --- lightly/models/modules/memory_bank.py | 5 +++-- tests/models/modules/test_memory_bank.py | 4 ++++ 2 files changed, 7 insertions(+), 2 deletions(-) diff --git a/lightly/models/modules/memory_bank.py b/lightly/models/modules/memory_bank.py index ae9b1eda4..63213ceba 100644 --- a/lightly/models/modules/memory_bank.py +++ b/lightly/models/modules/memory_bank.py @@ -24,7 +24,8 @@ class MemoryBankModule(Module): size: Size of the memory bank as (num_features, dim) tuple. If num_features is 0 then the memory bank is disabled. Pass 0 to disable the memory bank; any - other bare integer is rejected because the feature dimension is required. + other bare integer is rejected because the feature dimension is required + and cannot be guessed. gather_distributed: If True then negatives from all gpus are gathered before the memory bank is updated. This results in more frequent updates of the memory bank and @@ -53,7 +54,7 @@ class MemoryBankModule(Module): def __init__( self, - size: Union[int, Sequence[int]] = (65536, 128), + size: Union[int, Sequence[int]], gather_distributed: bool = False, feature_dim_first: bool = True, ): diff --git a/tests/models/modules/test_memory_bank.py b/tests/models/modules/test_memory_bank.py index d7e0e3e4f..ff0f74e9b 100644 --- a/tests/models/modules/test_memory_bank.py +++ b/tests/models/modules/test_memory_bank.py @@ -91,6 +91,10 @@ def test_init__no_dim(self) -> None: ): MemoryBankModule(size=10) + def test_init__size_required(self) -> None: + with pytest.raises(TypeError): + MemoryBankModule() # type: ignore[call-arg] + def test_init__disabled(self) -> None: memory_bank = MemoryBankModule(size=0) x = torch.randn(3, 2) From e07da832fccfe9c20db8516acf4a2f8912f37c6f Mon Sep 17 00:00:00 2001 From: gabrielfruet Date: Thu, 10 Sep 2026 18:34:03 -0300 Subject: [PATCH 8/9] docs: drop deleted NNCLR wrapper from NNMemoryBankModule example Co-Authored-By: Claude Opus 5 --- lightly/models/modules/nn_memory_bank.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/lightly/models/modules/nn_memory_bank.py b/lightly/models/modules/nn_memory_bank.py index a264ce161..7b56ecd0d 100644 --- a/lightly/models/modules/nn_memory_bank.py +++ b/lightly/models/modules/nn_memory_bank.py @@ -26,13 +26,12 @@ class NNMemoryBankModule(MemoryBankModule): then the memory bank is disabled. Examples: - >>> model = NNCLR(backbone) >>> criterion = NTXentLoss(temperature=0.1) >>> - >>> nn_replacer = NNmemoryBankModule(size=(2**16, 128)) + >>> nn_replacer = NNMemoryBankModule(size=(2**16, 128)) >>> >>> # forward pass - >>> (z0, p0), (z1, p1) = model(x0, x1) + >>> (z0, p0), (z1, p1) = model(x0), model(x1) >>> z0 = nn_replacer(z0.detach(), update=False) >>> z1 = nn_replacer(z1.detach(), update=True) >>> @@ -40,7 +39,7 @@ class NNMemoryBankModule(MemoryBankModule): """ - def __init__(self, size: Union[int, Sequence[int]] = (2**16, 128)): + def __init__(self, size: Union[int, Sequence[int]]): super(NNMemoryBankModule, self).__init__(size) def forward( # type: ignore[override] # TODO(Philipp, 11/23): Fix signature to match parent class. From 5ae62c77f2ee9089140c2242da3980256d187453 Mon Sep 17 00:00:00 2001 From: gabrielfruet Date: Fri, 11 Sep 2026 10:11:04 -0300 Subject: [PATCH 9/9] fix: reject memory bank sizes that are not (num_features, dim) Dropping the lazy init from forward() left size=(10,) building an uninitialized bank instead of raising. Co-Authored-By: Claude Opus 5 --- lightly/models/modules/memory_bank.py | 5 +++++ tests/models/modules/test_memory_bank.py | 12 ++++++++++++ 2 files changed, 17 insertions(+) diff --git a/lightly/models/modules/memory_bank.py b/lightly/models/modules/memory_bank.py index 63213ceba..d5041e2c9 100644 --- a/lightly/models/modules/memory_bank.py +++ b/lightly/models/modules/memory_bank.py @@ -71,6 +71,11 @@ def __init__( f"dimension. Set it with 'size=({size}, dim)', or pass 'size=0' to " "disable the memory bank." ) + if size_tuple != (0,) and len(size_tuple) != 2: + raise ValueError( + f"Illegal memory bank size {size}, expected a (num_features, dim) " + "tuple, or 'size=0' to disable the memory bank." + ) self.size = size_tuple self.gather_distributed = gather_distributed diff --git a/tests/models/modules/test_memory_bank.py b/tests/models/modules/test_memory_bank.py index ff0f74e9b..603250107 100644 --- a/tests/models/modules/test_memory_bank.py +++ b/tests/models/modules/test_memory_bank.py @@ -1,4 +1,5 @@ import re +from typing import Tuple import pytest import torch @@ -91,6 +92,17 @@ def test_init__no_dim(self) -> None: ): MemoryBankModule(size=10) + @pytest.mark.parametrize("size", [(), (10,), (10, 2, 3)]) + def test_init__invalid_shape(self, size: Tuple[int, ...]) -> None: + with pytest.raises( + ValueError, + match=re.escape( + f"Illegal memory bank size {size}, expected a (num_features, dim) " + "tuple, or 'size=0' to disable the memory bank." + ), + ): + MemoryBankModule(size=size) + def test_init__size_required(self) -> None: with pytest.raises(TypeError): MemoryBankModule() # type: ignore[call-arg]