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/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) 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/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/__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..2484adf92 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. @@ -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. @@ -71,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 0cee47d5d..ee2c898bf 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) @@ -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/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/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/__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/modules/memory_bank.py b/lightly/models/modules/memory_bank.py index 7c7606abd..d5041e2c9 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,9 @@ 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 + 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 @@ -57,7 +54,7 @@ class MemoryBankModule(Module): def __init__( self, - size: Union[int, Sequence[int]] = 65536, + size: Union[int, Sequence[int]], gather_distributed: bool = False, feature_dim_first: bool = True, ): @@ -68,6 +65,17 @@ 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." + ) + 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 @@ -85,17 +93,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 +123,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..7b56ecd0d 100644 --- a/lightly/models/modules/nn_memory_bank.py +++ b/lightly/models/modules/nn_memory_bank.py @@ -23,20 +23,15 @@ 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) >>> 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) >>> @@ -44,7 +39,7 @@ class NNMemoryBankModule(MemoryBankModule): """ - def __init__(self, size: Union[int, Sequence[int]] = 2**16): + 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. 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/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 4830c79aa..993dc59e9 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, cast -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 @@ -64,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 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/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/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 6d7e719d3..29060eeef 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -208,20 +208,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 @@ -270,7 +256,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 | @@ -290,7 +275,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 | @@ -303,20 +287,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/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/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/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..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 @@ -15,7 +16,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 +41,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 +54,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 +81,39 @@ 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) + @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] + + 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 +154,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) 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"]) 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] 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)