diff --git a/benchmarks/imagenet/vitb16/main.py b/benchmarks/imagenet/vitb16/main.py index 661f1cbfd..764b15cac 100644 --- a/benchmarks/imagenet/vitb16/main.py +++ b/benchmarks/imagenet/vitb16/main.py @@ -12,6 +12,7 @@ import lejepa import linear_eval import mae +import pixio import torch from pytorch_lightning import LightningModule, Trainer, seed_everything from pytorch_lightning.callbacks import ( @@ -58,6 +59,7 @@ "ibot": {"model": ibot.IBOT, "transform": ibot.transform}, "lejepa": {"model": lejepa.LeJEPA, "transform": lejepa.transform}, "mae": {"model": mae.MAE, "transform": mae.transform}, + "pixio": {"model": pixio.Pixio, "transform": pixio.transform}, "aim": {"model": aim.AIM, "transform": aim.transform}, } diff --git a/benchmarks/imagenet/vitb16/pixio.py b/benchmarks/imagenet/vitb16/pixio.py new file mode 100644 index 000000000..c27df9101 --- /dev/null +++ b/benchmarks/imagenet/vitb16/pixio.py @@ -0,0 +1,162 @@ +from typing import List, Tuple + +from pytorch_lightning import LightningModule +from timm.models.vision_transformer import vit_base_patch16_224 +from torch import Tensor +from torch.nn import MSELoss +from torch.optim import AdamW + +from lightly.models import utils +from lightly.models.modules import MaskedVisionTransformerTIMM, PixioDecoderTIMM +from lightly.transforms import MAETransform +from lightly.utils.benchmarking import OnlineLinearClassifier +from lightly.utils.scheduler import CosineWarmupScheduler + + +class Pixio(LightningModule): + def __init__(self, batch_size_per_device: int, num_classes: int) -> None: + super().__init__() + self.save_hyperparameters() + self.batch_size_per_device = batch_size_per_device + + decoder_dim = 512 + self.mask_ratio = 0.75 + self.grid_size = 4 + # vit-b/16 at 256px with 8 prefix tokens (1 cls + 7 reg). dynamic_img_size + # lets downstream evaluation run at other resolutions via pos-embed resampling. + vit = vit_base_patch16_224(img_size=256, reg_tokens=7, dynamic_img_size=True) + self.num_prefix_tokens = vit.num_prefix_tokens + self.patch_size = vit.patch_embed.patch_size[0] + self.sequence_length = vit.patch_embed.num_patches + vit.num_prefix_tokens + self.backbone = MaskedVisionTransformerTIMM(vit=vit) + self.decoder = PixioDecoderTIMM( + num_patches=vit.patch_embed.num_patches, + patch_size=self.patch_size, + embed_dim=vit.embed_dim, + decoder_embed_dim=decoder_dim, + decoder_depth=32, + decoder_num_heads=16, + num_prefix_tokens=self.num_prefix_tokens, + mlp_ratio=4.0, + proj_drop_rate=0.0, + attn_drop_rate=0.0, + ) + self.criterion = MSELoss() + + self.online_classifier = OnlineLinearClassifier( + feature_dim=vit.embed_dim, num_classes=num_classes + ) + + def forward(self, x: Tensor) -> Tensor: + # global representation = mean over the prefix (class) tokens + features = self.backbone.encode(images=x) + return features[:, : self.num_prefix_tokens].mean(dim=1) + + def forward_encoder(self, images: Tensor, idx_keep: Tensor) -> Tensor: + return self.backbone.encode(images=images, idx_keep=idx_keep) + + def forward_decoder( + self, x_encoded: Tensor, idx_keep: Tensor, idx_mask: Tensor + ) -> Tensor: + # build decoder input + batch_size = x_encoded.shape[0] + x_decode = self.decoder.embed(x_encoded) + x_masked = utils.repeat_token( + self.decoder.mask_token, (batch_size, self.sequence_length) + ) + x_masked = utils.set_at_index(x_masked, idx_keep, x_decode.type_as(x_masked)) + + # decoder forward pass + x_decoded = self.decoder.decode(x_masked) + + # predict pixel values for masked tokens + x_pred = utils.get_at_index(x_decoded, idx_mask) + x_pred = self.decoder.predict(x_pred) + return x_pred + + def training_step( + self, batch: Tuple[List[Tensor], Tensor, List[str]], batch_idx: int + ) -> Tensor: + images, targets = batch[0], batch[1] + images = images[0] # images is a list containing only one view + batch_size = images.shape[0] + idx_keep, idx_mask = utils.random_grid_token_mask( + size=(batch_size, self.sequence_length), + mask_ratio=self.mask_ratio, + grid_size=self.grid_size, + num_prefix_tokens=self.num_prefix_tokens, + device=images.device, + ) + x_encoded = self.forward_encoder(images=images, idx_keep=idx_keep) + predictions = self.forward_decoder( + x_encoded=x_encoded, idx_keep=idx_keep, idx_mask=idx_mask + ) + + # reconstruction target: normalized pixel values of the masked patches + patches = utils.patchify(images, self.patch_size) + target = utils.get_at_index(patches, idx_mask - self.num_prefix_tokens) + target = utils.normalize_mean_var(target) + + loss = self.criterion(predictions, target) + self.log( + "train_loss", loss, prog_bar=True, sync_dist=True, batch_size=len(targets) + ) + + cls_features = x_encoded[:, : self.num_prefix_tokens].mean(dim=1) + cls_loss, cls_log = self.online_classifier.training_step( + (cls_features.detach(), targets), batch_idx + ) + self.log_dict(cls_log, sync_dist=True, batch_size=len(targets)) + return loss + cls_loss + + def validation_step( + self, batch: Tuple[Tensor, Tensor, List[str]], batch_idx: int + ) -> Tensor: + images, targets = batch[0], batch[1] + cls_features = self.forward(images).flatten(start_dim=1) + cls_loss, cls_log = self.online_classifier.validation_step( + (cls_features.detach(), targets), batch_idx + ) + self.log_dict(cls_log, prog_bar=True, sync_dist=True, batch_size=len(targets)) + return cls_loss + + def configure_optimizers(self): + # Don't use weight decay for batch norm, bias parameters, and classification + # head to improve performance. + params, params_no_weight_decay = utils.get_weight_decay_parameters( + [self.backbone, self.decoder] + ) + optimizer = AdamW( + [ + {"name": "pixio", "params": params}, + { + "name": "pixio_no_weight_decay", + "params": params_no_weight_decay, + "weight_decay": 0.0, + }, + { + "name": "online_classifier", + "params": self.online_classifier.parameters(), + "weight_decay": 0.0, + }, + ], + lr=1.5e-4 * self.batch_size_per_device * self.trainer.world_size / 256, + weight_decay=0.05, + betas=(0.9, 0.95), + ) + scheduler = { + "scheduler": CosineWarmupScheduler( + optimizer=optimizer, + warmup_epochs=( + self.trainer.estimated_stepping_batches + / self.trainer.max_epochs + * 40 + ), + max_epochs=self.trainer.estimated_stepping_batches, + ), + "interval": "step", + } + return [optimizer], [scheduler] + + +transform = MAETransform(input_size=256) diff --git a/docs/source/examples/models.rst b/docs/source/examples/models.rst index a96938dc4..8f84d9d9e 100644 --- a/docs/source/examples/models.rst +++ b/docs/source/examples/models.rst @@ -27,6 +27,7 @@ for PyTorch and PyTorch Lightning to give you a headstart when implementing your msn.rst moco.rst nnclr.rst + pixio.rst pmsn.rst simclr.rst simmim.rst diff --git a/docs/source/examples/pixio.rst b/docs/source/examples/pixio.rst new file mode 100644 index 000000000..85eff7a88 --- /dev/null +++ b/docs/source/examples/pixio.rst @@ -0,0 +1,90 @@ +.. _pixio: + +Pixio +===== + +Example implementation of the Pixio method. Pixio builds on the `Masked Autoencoder +(MAE) `_ and adapts it for dense prediction through +three changes: a much deeper decoder (32 blocks) that takes over pixel-level detail +modeling, a larger masking granularity that masks whole blocks of patches on a regular +grid instead of individual patches, and multiple class tokens whose mean is used as the +global image representation. + +Key Components +-------------- + +- **Data Augmentations**: Like MAE, Pixio relies only on random resized cropping. +- **Masking**: Pixio masks 75% of the patches, but at a coarser granularity: whole + ``grid_size`` x ``grid_size`` blocks of patches are masked together (4x4 by default), + which prevents trivial reconstruction from neighboring patches. +- **Backbone**: A standard ViT with multiple class tokens (8 by default, realized via + ``reg_tokens``). +- **Decoder**: A deep (32-block) decoder that reconstructs the masked pixels. +- **Reconstruction Loss**: A Mean Squared Error (MSE) loss between the predicted and the + normalized pixel values of the masked patches. + +Good to Know +------------ + +- **Masking granularity**: The paper's headline configuration uses a 4x4 grid and 8 + class tokens. The dense-prediction-optimal ablation uses a 2x2 grid and 4 class + tokens. +- **Input resolution**: The reference model is trained at 256x256 with patch size 16 so + that the 16x16 patch grid divides evenly into 4x4 blocks. + +Reference: + `In Pursuit of Pixel Supervision for Visual Pre-training, 2025 `_ + +.. note:: + + Pixio requires `TIMM `_ to be + installed + + .. code-block:: bash + + pip install "lightly[timm]" + +.. tabs:: + .. tab:: PyTorch + + .. image:: https://img.shields.io/badge/Open%20in%20Colab-blue?logo=googlecolab&label=%20&labelColor=5c5c5c + :target: https://colab.research.google.com/github/lightly-ai/lightly/blob/master/examples/notebooks/pytorch/pixio.ipynb + + This example can be run from the command line with:: + + python lightly/examples/pytorch/pixio.py + + .. literalinclude:: ../../../examples/pytorch/pixio.py + + .. tab:: Lightning + + .. image:: https://img.shields.io/badge/Open%20in%20Colab-blue?logo=googlecolab&label=%20&labelColor=5c5c5c + :target: https://colab.research.google.com/github/lightly-ai/lightly/blob/master/examples/notebooks/pytorch_lightning/pixio.ipynb + + This example can be run from the command line with:: + + python lightly/examples/pytorch_lightning/pixio.py + + .. literalinclude:: ../../../examples/pytorch_lightning/pixio.py + + .. tab:: Lightning Distributed + + .. image:: https://img.shields.io/badge/Open%20in%20Colab-blue?logo=googlecolab&label=%20&labelColor=5c5c5c + :target: https://colab.research.google.com/github/lightly-ai/lightly/blob/master/examples/notebooks/pytorch_lightning_distributed/pixio.ipynb + + This example runs on multiple gpus using Distributed Data Parallel (DDP) + training with Pytorch Lightning. At least one GPU must be available on + the system. The example can be run from the command line with:: + + python lightly/examples/pytorch_lightning_distributed/pixio.py + + The model differs in the following ways from the non-distributed + implementation: + + - Distributed Data Parallel is enabled + - Distributed Sampling is used in the dataloader + + Distributed Sampling makes sure that each distributed process sees only + a subset of the data. + + .. literalinclude:: ../../../examples/pytorch_lightning_distributed/pixio.py diff --git a/examples/notebooks/pytorch/pixio.ipynb b/examples/notebooks/pytorch/pixio.ipynb new file mode 100644 index 000000000..d912df56b --- /dev/null +++ b/examples/notebooks/pytorch/pixio.ipynb @@ -0,0 +1,247 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "0", + "metadata": {}, + "source": [ + "This example requires the following dependencies to be installed:\n", + "pip install \"lightly[timm]\"" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "1", + "metadata": {}, + "outputs": [], + "source": [ + "!pip install \"lightly[timm]\"" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "2", + "metadata": {}, + "outputs": [], + "source": [ + "# Note: The model and training settings do not follow the reference settings\n", + "# from the paper. The settings are chosen such that the example can easily be\n", + "# run on a small dataset with a single GPU.\n", + "import torch\n", + "import torchvision\n", + "from timm.models.vision_transformer import vit_small_patch16_224\n", + "from torch import nn" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "3", + "metadata": {}, + "outputs": [], + "source": [ + "from lightly.models import utils\n", + "from lightly.models.modules import MaskedVisionTransformerTIMM, PixioDecoderTIMM\n", + "from lightly.transforms import MAETransform" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4", + "metadata": {}, + "outputs": [], + "source": [ + "class Pixio(nn.Module):\n", + " def __init__(self):\n", + " super().__init__()\n", + "\n", + " decoder_dim = 512\n", + " vit = vit_small_patch16_224(img_size=256, reg_tokens=7)\n", + " self.mask_ratio = 0.75\n", + " self.grid_size = 4\n", + " self.patch_size = vit.patch_embed.patch_size[0]\n", + " self.num_prefix_tokens = vit.num_prefix_tokens\n", + "\n", + " self.backbone = MaskedVisionTransformerTIMM(vit=vit)\n", + " self.sequence_length = self.backbone.sequence_length\n", + " self.decoder = PixioDecoderTIMM(\n", + " num_patches=vit.patch_embed.num_patches,\n", + " patch_size=self.patch_size,\n", + " embed_dim=vit.embed_dim,\n", + " decoder_embed_dim=decoder_dim,\n", + " decoder_depth=32,\n", + " decoder_num_heads=16,\n", + " num_prefix_tokens=self.num_prefix_tokens,\n", + " mlp_ratio=4.0,\n", + " proj_drop_rate=0.0,\n", + " attn_drop_rate=0.0,\n", + " )\n", + "\n", + " def forward_encoder(self, images, idx_keep=None):\n", + " return self.backbone.encode(images=images, idx_keep=idx_keep)\n", + "\n", + " def forward_decoder(self, x_encoded, idx_keep, idx_mask):\n", + " # build decoder input\n", + " batch_size = x_encoded.shape[0]\n", + " x_decode = self.decoder.embed(x_encoded)\n", + " x_masked = utils.repeat_token(\n", + " self.decoder.mask_token, (batch_size, self.sequence_length)\n", + " )\n", + " x_masked = utils.set_at_index(x_masked, idx_keep, x_decode.type_as(x_masked))\n", + "\n", + " # decoder forward pass\n", + " x_decoded = self.decoder.decode(x_masked)\n", + "\n", + " # predict pixel values for masked tokens\n", + " x_pred = utils.get_at_index(x_decoded, idx_mask)\n", + " x_pred = self.decoder.predict(x_pred)\n", + " return x_pred\n", + "\n", + " def forward(self, images):\n", + " batch_size = images.shape[0]\n", + " idx_keep, idx_mask = utils.random_grid_token_mask(\n", + " size=(batch_size, self.sequence_length),\n", + " mask_ratio=self.mask_ratio,\n", + " grid_size=self.grid_size,\n", + " num_prefix_tokens=self.num_prefix_tokens,\n", + " device=images.device,\n", + " )\n", + " x_encoded = self.forward_encoder(images=images, idx_keep=idx_keep)\n", + " x_pred = self.forward_decoder(\n", + " x_encoded=x_encoded, idx_keep=idx_keep, idx_mask=idx_mask\n", + " )\n", + "\n", + " # get image patches for masked tokens\n", + " patches = utils.patchify(images, self.patch_size)\n", + " # must adjust idx_mask for the prefix tokens\n", + " target = utils.get_at_index(patches, idx_mask - self.num_prefix_tokens)\n", + " target = utils.normalize_mean_var(target)\n", + " return x_pred, target" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "5", + "metadata": {}, + "outputs": [], + "source": [ + "model = Pixio()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "6", + "metadata": {}, + "outputs": [], + "source": [ + "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n", + "model.to(device)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "7", + "metadata": { + "lines_to_next_cell": 2 + }, + "outputs": [], + "source": [ + "transform = MAETransform(input_size=256)\n", + "# we ignore object detection annotations by setting target_transform to return 0" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "8", + "metadata": {}, + "outputs": [], + "source": [ + "def target_transform(t):\n", + " return 0" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "9", + "metadata": {}, + "outputs": [], + "source": [ + "dataset = torchvision.datasets.VOCDetection(\n", + " \"datasets/pascal_voc\",\n", + " download=True,\n", + " transform=transform,\n", + " target_transform=target_transform,\n", + ")\n", + "# or create a dataset from a folder containing images or videos:\n", + "# dataset = LightlyDataset(\"path/to/folder\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "10", + "metadata": {}, + "outputs": [], + "source": [ + "dataloader = torch.utils.data.DataLoader(\n", + " dataset,\n", + " batch_size=64,\n", + " shuffle=True,\n", + " drop_last=True,\n", + " num_workers=8,\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "11", + "metadata": {}, + "outputs": [], + "source": [ + "criterion = nn.MSELoss()\n", + "optimizer = torch.optim.AdamW(model.parameters(), lr=1.5e-4)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "12", + "metadata": {}, + "outputs": [], + "source": [ + "print(\"Starting Training\")\n", + "for epoch in range(10):\n", + " total_loss = 0\n", + " for batch in dataloader:\n", + " views = batch[0]\n", + " images = views[0].to(device) # views contains only a single view\n", + " predictions, targets = model(images)\n", + " loss = criterion(predictions, targets)\n", + " total_loss += loss.detach()\n", + " loss.backward()\n", + " optimizer.step()\n", + " optimizer.zero_grad()\n", + " avg_loss = total_loss / len(dataloader)\n", + " print(f\"epoch: {epoch:>02}, loss: {avg_loss:.5f}\")" + ] + } + ], + "metadata": { + "jupytext": { + "cell_metadata_filter": "-all", + "main_language": "python", + "notebook_metadata_filter": "-all" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/examples/notebooks/pytorch_lightning/pixio.ipynb b/examples/notebooks/pytorch_lightning/pixio.ipynb new file mode 100644 index 000000000..376012a5a --- /dev/null +++ b/examples/notebooks/pytorch_lightning/pixio.ipynb @@ -0,0 +1,234 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "0", + "metadata": {}, + "source": [ + "This example requires the following dependencies to be installed:\n", + "pip install \"lightly[timm]\"" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "1", + "metadata": {}, + "outputs": [], + "source": [ + "!pip install \"lightly[timm]\"" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "2", + "metadata": {}, + "outputs": [], + "source": [ + "# Note: The model and training settings do not follow the reference settings\n", + "# from the paper. The settings are chosen such that the example can easily be\n", + "# run on a small dataset with a single GPU.\n", + "import pytorch_lightning as pl\n", + "import torch\n", + "import torchvision\n", + "from timm.models.vision_transformer import vit_small_patch16_224\n", + "from torch import nn" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "3", + "metadata": {}, + "outputs": [], + "source": [ + "from lightly.models import utils\n", + "from lightly.models.modules import MaskedVisionTransformerTIMM, PixioDecoderTIMM\n", + "from lightly.transforms import MAETransform" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4", + "metadata": {}, + "outputs": [], + "source": [ + "class Pixio(pl.LightningModule):\n", + " def __init__(self):\n", + " super().__init__()\n", + "\n", + " decoder_dim = 512\n", + " vit = vit_small_patch16_224(img_size=256, reg_tokens=7)\n", + " self.mask_ratio = 0.75\n", + " self.grid_size = 4\n", + " self.patch_size = vit.patch_embed.patch_size[0]\n", + " self.num_prefix_tokens = vit.num_prefix_tokens\n", + " self.backbone = MaskedVisionTransformerTIMM(vit=vit)\n", + " self.sequence_length = self.backbone.sequence_length\n", + " self.decoder = PixioDecoderTIMM(\n", + " num_patches=vit.patch_embed.num_patches,\n", + " patch_size=self.patch_size,\n", + " embed_dim=vit.embed_dim,\n", + " decoder_embed_dim=decoder_dim,\n", + " decoder_depth=32,\n", + " decoder_num_heads=16,\n", + " num_prefix_tokens=self.num_prefix_tokens,\n", + " mlp_ratio=4.0,\n", + " proj_drop_rate=0.0,\n", + " attn_drop_rate=0.0,\n", + " )\n", + " self.criterion = nn.MSELoss()\n", + "\n", + " def forward_encoder(self, images, idx_keep=None):\n", + " return self.backbone.encode(images=images, idx_keep=idx_keep)\n", + "\n", + " def forward_decoder(self, x_encoded, idx_keep, idx_mask):\n", + " # build decoder input\n", + " batch_size = x_encoded.shape[0]\n", + " x_decode = self.decoder.embed(x_encoded)\n", + " x_masked = utils.repeat_token(\n", + " self.decoder.mask_token, (batch_size, self.sequence_length)\n", + " )\n", + " x_masked = utils.set_at_index(x_masked, idx_keep, x_decode.type_as(x_masked))\n", + "\n", + " # decoder forward pass\n", + " x_decoded = self.decoder.decode(x_masked)\n", + "\n", + " # predict pixel values for masked tokens\n", + " x_pred = utils.get_at_index(x_decoded, idx_mask)\n", + " x_pred = self.decoder.predict(x_pred)\n", + " return x_pred\n", + "\n", + " def training_step(self, batch, batch_idx):\n", + " views = batch[0]\n", + " images = views[0] # views contains only a single view\n", + " batch_size = images.shape[0]\n", + " idx_keep, idx_mask = utils.random_grid_token_mask(\n", + " size=(batch_size, self.sequence_length),\n", + " mask_ratio=self.mask_ratio,\n", + " grid_size=self.grid_size,\n", + " num_prefix_tokens=self.num_prefix_tokens,\n", + " device=images.device,\n", + " )\n", + " x_encoded = self.forward_encoder(images=images, idx_keep=idx_keep)\n", + " x_pred = self.forward_decoder(\n", + " x_encoded=x_encoded, idx_keep=idx_keep, idx_mask=idx_mask\n", + " )\n", + "\n", + " # get image patches for masked tokens\n", + " patches = utils.patchify(images, self.patch_size)\n", + " # must adjust idx_mask for the prefix tokens\n", + " target = utils.get_at_index(patches, idx_mask - self.num_prefix_tokens)\n", + " target = utils.normalize_mean_var(target)\n", + "\n", + " loss = self.criterion(x_pred, target)\n", + " return loss\n", + "\n", + " def configure_optimizers(self):\n", + " optim = torch.optim.AdamW(self.parameters(), lr=1.5e-4)\n", + " return optim" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "5", + "metadata": {}, + "outputs": [], + "source": [ + "model = Pixio()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "6", + "metadata": {}, + "outputs": [], + "source": [ + "transform = MAETransform(input_size=256)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "7", + "metadata": {}, + "outputs": [], + "source": [ + "# we ignore object detection annotations by setting target_transform to return 0\n", + "def target_transform(t):\n", + " return 0" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "8", + "metadata": {}, + "outputs": [], + "source": [ + "dataset = torchvision.datasets.VOCDetection(\n", + " \"datasets/pascal_voc\",\n", + " download=True,\n", + " transform=transform,\n", + " target_transform=target_transform,\n", + ")\n", + "# or create a dataset from a folder containing images or videos:\n", + "# dataset = LightlyDataset(\"path/to/folder\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "9", + "metadata": {}, + "outputs": [], + "source": [ + "dataloader = torch.utils.data.DataLoader(\n", + " dataset,\n", + " batch_size=64,\n", + " shuffle=True,\n", + " drop_last=True,\n", + " num_workers=8,\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "10", + "metadata": {}, + "outputs": [], + "source": [ + "accelerator = \"gpu\" if torch.cuda.is_available() else \"cpu\"" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "11", + "metadata": {}, + "outputs": [], + "source": [ + "trainer = pl.Trainer(\n", + " max_epochs=10,\n", + " devices=1,\n", + " accelerator=accelerator,\n", + ")\n", + "trainer.fit(model=model, train_dataloaders=dataloader)" + ] + } + ], + "metadata": { + "jupytext": { + "cell_metadata_filter": "-all", + "main_language": "python", + "notebook_metadata_filter": "-all" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/examples/notebooks/pytorch_lightning_distributed/pixio.ipynb b/examples/notebooks/pytorch_lightning_distributed/pixio.ipynb new file mode 100644 index 000000000..06bbd443a --- /dev/null +++ b/examples/notebooks/pytorch_lightning_distributed/pixio.ipynb @@ -0,0 +1,228 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "0", + "metadata": {}, + "source": [ + "This example requires the following dependencies to be installed:\n", + "pip install \"lightly[timm]\"" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "1", + "metadata": {}, + "outputs": [], + "source": [ + "!pip install \"lightly[timm]\"" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "2", + "metadata": {}, + "outputs": [], + "source": [ + "# Note: The model and training settings do not follow the reference settings\n", + "# from the paper. The settings are chosen such that the example can easily be\n", + "# run on a small dataset with a single GPU.\n", + "import pytorch_lightning as pl\n", + "import torch\n", + "import torchvision\n", + "from timm.models.vision_transformer import vit_small_patch16_224\n", + "from torch import nn" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "3", + "metadata": {}, + "outputs": [], + "source": [ + "from lightly.models import utils\n", + "from lightly.models.modules import MaskedVisionTransformerTIMM, PixioDecoderTIMM\n", + "from lightly.transforms import MAETransform" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4", + "metadata": {}, + "outputs": [], + "source": [ + "class Pixio(pl.LightningModule):\n", + " def __init__(self):\n", + " super().__init__()\n", + "\n", + " decoder_dim = 512\n", + " vit = vit_small_patch16_224(img_size=256, reg_tokens=7)\n", + " self.mask_ratio = 0.75\n", + " self.grid_size = 4\n", + " self.patch_size = vit.patch_embed.patch_size[0]\n", + " self.num_prefix_tokens = vit.num_prefix_tokens\n", + " self.backbone = MaskedVisionTransformerTIMM(vit=vit)\n", + " self.sequence_length = self.backbone.sequence_length\n", + " self.decoder = PixioDecoderTIMM(\n", + " num_patches=vit.patch_embed.num_patches,\n", + " patch_size=self.patch_size,\n", + " embed_dim=vit.embed_dim,\n", + " decoder_embed_dim=decoder_dim,\n", + " decoder_depth=32,\n", + " decoder_num_heads=16,\n", + " num_prefix_tokens=self.num_prefix_tokens,\n", + " mlp_ratio=4.0,\n", + " proj_drop_rate=0.0,\n", + " attn_drop_rate=0.0,\n", + " )\n", + " self.criterion = nn.MSELoss()\n", + "\n", + " def forward_encoder(self, images, idx_keep=None):\n", + " return self.backbone.encode(images=images, idx_keep=idx_keep)\n", + "\n", + " def forward_decoder(self, x_encoded, idx_keep, idx_mask):\n", + " # build decoder input\n", + " batch_size = x_encoded.shape[0]\n", + " x_decode = self.decoder.embed(x_encoded)\n", + " x_masked = utils.repeat_token(\n", + " self.decoder.mask_token, (batch_size, self.sequence_length)\n", + " )\n", + " x_masked = utils.set_at_index(x_masked, idx_keep, x_decode.type_as(x_masked))\n", + "\n", + " # decoder forward pass\n", + " x_decoded = self.decoder.decode(x_masked)\n", + "\n", + " # predict pixel values for masked tokens\n", + " x_pred = utils.get_at_index(x_decoded, idx_mask)\n", + " x_pred = self.decoder.predict(x_pred)\n", + " return x_pred\n", + "\n", + " def training_step(self, batch, batch_idx):\n", + " views = batch[0]\n", + " images = views[0] # views contains only a single view\n", + " batch_size = images.shape[0]\n", + " idx_keep, idx_mask = utils.random_grid_token_mask(\n", + " size=(batch_size, self.sequence_length),\n", + " mask_ratio=self.mask_ratio,\n", + " grid_size=self.grid_size,\n", + " num_prefix_tokens=self.num_prefix_tokens,\n", + " device=images.device,\n", + " )\n", + " x_encoded = self.forward_encoder(images=images, idx_keep=idx_keep)\n", + " x_pred = self.forward_decoder(\n", + " x_encoded=x_encoded, idx_keep=idx_keep, idx_mask=idx_mask\n", + " )\n", + "\n", + " # get image patches for masked tokens\n", + " patches = utils.patchify(images, self.patch_size)\n", + " # must adjust idx_mask for the prefix tokens\n", + " target = utils.get_at_index(patches, idx_mask - self.num_prefix_tokens)\n", + " target = utils.normalize_mean_var(target)\n", + "\n", + " loss = self.criterion(x_pred, target)\n", + " return loss\n", + "\n", + " def configure_optimizers(self):\n", + " optim = torch.optim.AdamW(self.parameters(), lr=1.5e-4)\n", + " return optim" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "5", + "metadata": {}, + "outputs": [], + "source": [ + "model = Pixio()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "6", + "metadata": {}, + "outputs": [], + "source": [ + "transform = MAETransform(input_size=256)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "7", + "metadata": {}, + "outputs": [], + "source": [ + "# we ignore object detection annotations by setting target_transform to return 0\n", + "def target_transform(t):\n", + " return 0" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "8", + "metadata": {}, + "outputs": [], + "source": [ + "dataset = torchvision.datasets.VOCDetection(\n", + " \"datasets/pascal_voc\",\n", + " download=True,\n", + " transform=transform,\n", + " target_transform=target_transform,\n", + ")\n", + "# or create a dataset from a folder containing images or videos:\n", + "# dataset = LightlyDataset(\"path/to/folder\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "9", + "metadata": {}, + "outputs": [], + "source": [ + "dataloader = torch.utils.data.DataLoader(\n", + " dataset,\n", + " batch_size=64,\n", + " shuffle=True,\n", + " drop_last=True,\n", + " num_workers=8,\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "10", + "metadata": {}, + "outputs": [], + "source": [ + "# Train with DDP on multiple gpus. Distributed sampling is also enabled with\n", + "# replace_sampler_ddp=True.\n", + "trainer = pl.Trainer(\n", + " max_epochs=10,\n", + " devices=\"auto\",\n", + " accelerator=\"gpu\",\n", + " strategy=\"ddp_find_unused_parameters_true\",\n", + " use_distributed_sampler=True, # or replace_sampler_ddp=True for PyTorch Lightning <2.0\n", + ")\n", + "trainer.fit(model=model, train_dataloaders=dataloader)" + ] + } + ], + "metadata": { + "jupytext": { + "cell_metadata_filter": "-all", + "main_language": "python", + "notebook_metadata_filter": "-all" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/examples/pytorch/pixio.py b/examples/pytorch/pixio.py new file mode 100644 index 000000000..7d83d05c1 --- /dev/null +++ b/examples/pytorch/pixio.py @@ -0,0 +1,131 @@ +# This example requires the following dependencies to be installed: +# pip install "lightly[timm]" + +# Note: The model and training settings do not follow the reference settings +# from the paper. The settings are chosen such that the example can easily be +# run on a small dataset with a single GPU. +import torch +import torchvision +from timm.models.vision_transformer import vit_small_patch16_224 +from torch import nn + +from lightly.models import utils +from lightly.models.modules import MaskedVisionTransformerTIMM, PixioDecoderTIMM +from lightly.transforms import MAETransform + + +class Pixio(nn.Module): + def __init__(self): + super().__init__() + + decoder_dim = 512 + vit = vit_small_patch16_224(img_size=256, reg_tokens=7) + self.mask_ratio = 0.75 + self.grid_size = 4 + self.patch_size = vit.patch_embed.patch_size[0] + self.num_prefix_tokens = vit.num_prefix_tokens + + self.backbone = MaskedVisionTransformerTIMM(vit=vit) + self.sequence_length = self.backbone.sequence_length + self.decoder = PixioDecoderTIMM( + num_patches=vit.patch_embed.num_patches, + patch_size=self.patch_size, + embed_dim=vit.embed_dim, + decoder_embed_dim=decoder_dim, + decoder_depth=32, + decoder_num_heads=16, + num_prefix_tokens=self.num_prefix_tokens, + mlp_ratio=4.0, + proj_drop_rate=0.0, + attn_drop_rate=0.0, + ) + + def forward_encoder(self, images, idx_keep=None): + return self.backbone.encode(images=images, idx_keep=idx_keep) + + def forward_decoder(self, x_encoded, idx_keep, idx_mask): + # build decoder input + batch_size = x_encoded.shape[0] + x_decode = self.decoder.embed(x_encoded) + x_masked = utils.repeat_token( + self.decoder.mask_token, (batch_size, self.sequence_length) + ) + x_masked = utils.set_at_index(x_masked, idx_keep, x_decode.type_as(x_masked)) + + # decoder forward pass + x_decoded = self.decoder.decode(x_masked) + + # predict pixel values for masked tokens + x_pred = utils.get_at_index(x_decoded, idx_mask) + x_pred = self.decoder.predict(x_pred) + return x_pred + + def forward(self, images): + batch_size = images.shape[0] + idx_keep, idx_mask = utils.random_grid_token_mask( + size=(batch_size, self.sequence_length), + mask_ratio=self.mask_ratio, + grid_size=self.grid_size, + num_prefix_tokens=self.num_prefix_tokens, + device=images.device, + ) + x_encoded = self.forward_encoder(images=images, idx_keep=idx_keep) + x_pred = self.forward_decoder( + x_encoded=x_encoded, idx_keep=idx_keep, idx_mask=idx_mask + ) + + # get image patches for masked tokens + patches = utils.patchify(images, self.patch_size) + # must adjust idx_mask for the prefix tokens + target = utils.get_at_index(patches, idx_mask - self.num_prefix_tokens) + target = utils.normalize_mean_var(target) + return x_pred, target + + +model = Pixio() + +device = "cuda" if torch.cuda.is_available() else "cpu" +model.to(device) + +transform = MAETransform(input_size=256) +# we ignore object detection annotations by setting target_transform to return 0 + + +def target_transform(t): + return 0 + + +dataset = torchvision.datasets.VOCDetection( + "datasets/pascal_voc", + download=True, + transform=transform, + target_transform=target_transform, +) +# or create a dataset from a folder containing images or videos: +# dataset = LightlyDataset("path/to/folder") + +dataloader = torch.utils.data.DataLoader( + dataset, + batch_size=64, + shuffle=True, + drop_last=True, + num_workers=8, +) + +criterion = nn.MSELoss() +optimizer = torch.optim.AdamW(model.parameters(), lr=1.5e-4) + +print("Starting Training") +for epoch in range(10): + total_loss = 0 + for batch in dataloader: + views = batch[0] + images = views[0].to(device) # views contains only a single view + predictions, targets = model(images) + loss = criterion(predictions, targets) + total_loss += loss.detach() + loss.backward() + optimizer.step() + optimizer.zero_grad() + avg_loss = total_loss / len(dataloader) + print(f"epoch: {epoch:>02}, loss: {avg_loss:.5f}") diff --git a/examples/pytorch_lightning/pixio.py b/examples/pytorch_lightning/pixio.py new file mode 100644 index 000000000..93a56f67c --- /dev/null +++ b/examples/pytorch_lightning/pixio.py @@ -0,0 +1,128 @@ +# This example requires the following dependencies to be installed: +# pip install "lightly[timm]" + +# Note: The model and training settings do not follow the reference settings +# from the paper. The settings are chosen such that the example can easily be +# run on a small dataset with a single GPU. +import pytorch_lightning as pl +import torch +import torchvision +from timm.models.vision_transformer import vit_small_patch16_224 +from torch import nn + +from lightly.models import utils +from lightly.models.modules import MaskedVisionTransformerTIMM, PixioDecoderTIMM +from lightly.transforms import MAETransform + + +class Pixio(pl.LightningModule): + def __init__(self): + super().__init__() + + decoder_dim = 512 + vit = vit_small_patch16_224(img_size=256, reg_tokens=7) + self.mask_ratio = 0.75 + self.grid_size = 4 + self.patch_size = vit.patch_embed.patch_size[0] + self.num_prefix_tokens = vit.num_prefix_tokens + self.backbone = MaskedVisionTransformerTIMM(vit=vit) + self.sequence_length = self.backbone.sequence_length + self.decoder = PixioDecoderTIMM( + num_patches=vit.patch_embed.num_patches, + patch_size=self.patch_size, + embed_dim=vit.embed_dim, + decoder_embed_dim=decoder_dim, + decoder_depth=32, + decoder_num_heads=16, + num_prefix_tokens=self.num_prefix_tokens, + mlp_ratio=4.0, + proj_drop_rate=0.0, + attn_drop_rate=0.0, + ) + self.criterion = nn.MSELoss() + + def forward_encoder(self, images, idx_keep=None): + return self.backbone.encode(images=images, idx_keep=idx_keep) + + def forward_decoder(self, x_encoded, idx_keep, idx_mask): + # build decoder input + batch_size = x_encoded.shape[0] + x_decode = self.decoder.embed(x_encoded) + x_masked = utils.repeat_token( + self.decoder.mask_token, (batch_size, self.sequence_length) + ) + x_masked = utils.set_at_index(x_masked, idx_keep, x_decode.type_as(x_masked)) + + # decoder forward pass + x_decoded = self.decoder.decode(x_masked) + + # predict pixel values for masked tokens + x_pred = utils.get_at_index(x_decoded, idx_mask) + x_pred = self.decoder.predict(x_pred) + return x_pred + + def training_step(self, batch, batch_idx): + views = batch[0] + images = views[0] # views contains only a single view + batch_size = images.shape[0] + idx_keep, idx_mask = utils.random_grid_token_mask( + size=(batch_size, self.sequence_length), + mask_ratio=self.mask_ratio, + grid_size=self.grid_size, + num_prefix_tokens=self.num_prefix_tokens, + device=images.device, + ) + x_encoded = self.forward_encoder(images=images, idx_keep=idx_keep) + x_pred = self.forward_decoder( + x_encoded=x_encoded, idx_keep=idx_keep, idx_mask=idx_mask + ) + + # get image patches for masked tokens + patches = utils.patchify(images, self.patch_size) + # must adjust idx_mask for the prefix tokens + target = utils.get_at_index(patches, idx_mask - self.num_prefix_tokens) + target = utils.normalize_mean_var(target) + + loss = self.criterion(x_pred, target) + return loss + + def configure_optimizers(self): + optim = torch.optim.AdamW(self.parameters(), lr=1.5e-4) + return optim + + +model = Pixio() + +transform = MAETransform(input_size=256) + + +# we ignore object detection annotations by setting target_transform to return 0 +def target_transform(t): + return 0 + + +dataset = torchvision.datasets.VOCDetection( + "datasets/pascal_voc", + download=True, + transform=transform, + target_transform=target_transform, +) +# or create a dataset from a folder containing images or videos: +# dataset = LightlyDataset("path/to/folder") + +dataloader = torch.utils.data.DataLoader( + dataset, + batch_size=64, + shuffle=True, + drop_last=True, + num_workers=8, +) + +accelerator = "gpu" if torch.cuda.is_available() else "cpu" + +trainer = pl.Trainer( + max_epochs=10, + devices=1, + accelerator=accelerator, +) +trainer.fit(model=model, train_dataloaders=dataloader) diff --git a/examples/pytorch_lightning_distributed/pixio.py b/examples/pytorch_lightning_distributed/pixio.py new file mode 100644 index 000000000..f6dde4c65 --- /dev/null +++ b/examples/pytorch_lightning_distributed/pixio.py @@ -0,0 +1,130 @@ +# This example requires the following dependencies to be installed: +# pip install "lightly[timm]" + +# Note: The model and training settings do not follow the reference settings +# from the paper. The settings are chosen such that the example can easily be +# run on a small dataset with a single GPU. +import pytorch_lightning as pl +import torch +import torchvision +from timm.models.vision_transformer import vit_small_patch16_224 +from torch import nn + +from lightly.models import utils +from lightly.models.modules import MaskedVisionTransformerTIMM, PixioDecoderTIMM +from lightly.transforms import MAETransform + + +class Pixio(pl.LightningModule): + def __init__(self): + super().__init__() + + decoder_dim = 512 + vit = vit_small_patch16_224(img_size=256, reg_tokens=7) + self.mask_ratio = 0.75 + self.grid_size = 4 + self.patch_size = vit.patch_embed.patch_size[0] + self.num_prefix_tokens = vit.num_prefix_tokens + self.backbone = MaskedVisionTransformerTIMM(vit=vit) + self.sequence_length = self.backbone.sequence_length + self.decoder = PixioDecoderTIMM( + num_patches=vit.patch_embed.num_patches, + patch_size=self.patch_size, + embed_dim=vit.embed_dim, + decoder_embed_dim=decoder_dim, + decoder_depth=32, + decoder_num_heads=16, + num_prefix_tokens=self.num_prefix_tokens, + mlp_ratio=4.0, + proj_drop_rate=0.0, + attn_drop_rate=0.0, + ) + self.criterion = nn.MSELoss() + + def forward_encoder(self, images, idx_keep=None): + return self.backbone.encode(images=images, idx_keep=idx_keep) + + def forward_decoder(self, x_encoded, idx_keep, idx_mask): + # build decoder input + batch_size = x_encoded.shape[0] + x_decode = self.decoder.embed(x_encoded) + x_masked = utils.repeat_token( + self.decoder.mask_token, (batch_size, self.sequence_length) + ) + x_masked = utils.set_at_index(x_masked, idx_keep, x_decode.type_as(x_masked)) + + # decoder forward pass + x_decoded = self.decoder.decode(x_masked) + + # predict pixel values for masked tokens + x_pred = utils.get_at_index(x_decoded, idx_mask) + x_pred = self.decoder.predict(x_pred) + return x_pred + + def training_step(self, batch, batch_idx): + views = batch[0] + images = views[0] # views contains only a single view + batch_size = images.shape[0] + idx_keep, idx_mask = utils.random_grid_token_mask( + size=(batch_size, self.sequence_length), + mask_ratio=self.mask_ratio, + grid_size=self.grid_size, + num_prefix_tokens=self.num_prefix_tokens, + device=images.device, + ) + x_encoded = self.forward_encoder(images=images, idx_keep=idx_keep) + x_pred = self.forward_decoder( + x_encoded=x_encoded, idx_keep=idx_keep, idx_mask=idx_mask + ) + + # get image patches for masked tokens + patches = utils.patchify(images, self.patch_size) + # must adjust idx_mask for the prefix tokens + target = utils.get_at_index(patches, idx_mask - self.num_prefix_tokens) + target = utils.normalize_mean_var(target) + + loss = self.criterion(x_pred, target) + return loss + + def configure_optimizers(self): + optim = torch.optim.AdamW(self.parameters(), lr=1.5e-4) + return optim + + +model = Pixio() + +transform = MAETransform(input_size=256) + + +# we ignore object detection annotations by setting target_transform to return 0 +def target_transform(t): + return 0 + + +dataset = torchvision.datasets.VOCDetection( + "datasets/pascal_voc", + download=True, + transform=transform, + target_transform=target_transform, +) +# or create a dataset from a folder containing images or videos: +# dataset = LightlyDataset("path/to/folder") + +dataloader = torch.utils.data.DataLoader( + dataset, + batch_size=64, + shuffle=True, + drop_last=True, + num_workers=8, +) + +# Train with DDP on multiple gpus. Distributed sampling is also enabled with +# replace_sampler_ddp=True. +trainer = pl.Trainer( + max_epochs=10, + devices="auto", + accelerator="gpu", + strategy="ddp_find_unused_parameters_true", + use_distributed_sampler=True, # or replace_sampler_ddp=True for PyTorch Lightning <2.0 +) +trainer.fit(model=model, train_dataloaders=dataloader) diff --git a/lightly/models/modules/__init__.py b/lightly/models/modules/__init__.py index 57341294d..3e4bdb174 100644 --- a/lightly/models/modules/__init__.py +++ b/lightly/models/modules/__init__.py @@ -47,7 +47,10 @@ # Requires timm >= 0.9.9 from lightly.models.modules.heads_timm import AIMPredictionHead from lightly.models.modules.ijepa_timm import IJEPAPredictorTIMM - from lightly.models.modules.masked_autoencoder_timm import MAEDecoderTIMM + from lightly.models.modules.masked_autoencoder_timm import ( + MAEDecoderTIMM, + PixioDecoderTIMM, + ) from lightly.models.modules.masked_causal_vision_transformer import ( MaskedCausalVisionTransformer, ) diff --git a/lightly/models/modules/masked_autoencoder_timm.py b/lightly/models/modules/masked_autoencoder_timm.py index 7488e3150..8328a5f11 100644 --- a/lightly/models/modules/masked_autoencoder_timm.py +++ b/lightly/models/modules/masked_autoencoder_timm.py @@ -35,6 +35,9 @@ class MAEDecoderTIMM(Module): Depth of transformer. decoder_num_heads: Number of attention heads. + num_prefix_tokens: + Number of prefix tokens (e.g. class or register tokens) preceding the + patch tokens. Determines the length of the decoder positional embedding. mlp_ratio: Ratio of mlp hidden dim to embedding dim. proj_drop_rate: @@ -59,6 +62,7 @@ def __init__( decoder_embed_dim: int = 512, decoder_depth: int = 8, decoder_num_heads: int = 16, + num_prefix_tokens: int = 1, mlp_ratio: float = 4.0, proj_drop_rate: float = 0.0, attn_drop_rate: float = 0.0, @@ -69,6 +73,7 @@ def __init__( """Initializes the MAEDecoderTIMM with the specified parameters.""" super().__init__() + self.num_prefix_tokens = num_prefix_tokens self.decoder_embed = nn.Linear(embed_dim, decoder_embed_dim, bias=True) self.mask_token = ( nn.Parameter(torch.zeros(1, 1, decoder_embed_dim)) @@ -78,7 +83,8 @@ def __init__( # Positional encoding of the decoder self.decoder_pos_embed = nn.Parameter( - torch.zeros(1, num_patches + 1, decoder_embed_dim), requires_grad=False + torch.zeros(1, num_patches + num_prefix_tokens, decoder_embed_dim), + requires_grad=False, ) # fixed sin-cos embedding self.decoder_blocks = Sequential( @@ -176,6 +182,87 @@ def _initialize_weights(self) -> None: """Initializes weights for the decoder components.""" torch.nn.init.normal_(self.mask_token, std=0.02) utils.initialize_2d_sine_cosine_positional_embedding( - pos_embedding=self.decoder_pos_embed, has_class_token=True + pos_embedding=self.decoder_pos_embed, + num_prefix_tokens=self.num_prefix_tokens, ) self.apply(init_weights) + + +class PixioDecoderTIMM(MAEDecoderTIMM): + """Decoder for the Pixio model [0]. + + Pixio is a masked autoencoder variant that uses a much deeper decoder than MAE + (32 blocks vs. 8) to move pixel-level detail modeling from the encoder into the + decoder. This is an :class:`MAEDecoderTIMM` with a deeper default decoder + (``decoder_depth`` defaults to 32 following Pixio). Implemented from the paper; + not derived from the reference code. + + - [0]: In Pursuit of Pixel Supervision for Visual Pre-training, 2025, + https://arxiv.org/abs/2512.15715 + + Attributes: + num_patches: + Number of patches. + patch_size: + Patch size. + in_chans: + Number of image input channels. + embed_dim: + Embedding dimension of the encoder. + decoder_embed_dim: + Embedding dimension of the decoder. + decoder_depth: + Depth of the decoder transformer. Defaults to 32 following Pixio. + decoder_num_heads: + Number of attention heads. + num_prefix_tokens: + Number of prefix (class or register) tokens preceding the patch tokens. + mlp_ratio: + Ratio of mlp hidden dim to embedding dim. + proj_drop_rate: + Percentage of elements set to zero after the MLP in the transformer. + attn_drop_rate: + Percentage of elements set to zero after the attention head. + norm_layer: + Normalization layer. + initialize_weights: + Flag that determines if weights should be initialized. + mask_token: + The mask token. + + """ + + def __init__( + self, + num_patches: int, + patch_size: int, + in_chans: int = 3, + embed_dim: int = 1024, + decoder_embed_dim: int = 512, + decoder_depth: int = 32, + decoder_num_heads: int = 16, + num_prefix_tokens: int = 1, + mlp_ratio: float = 4.0, + proj_drop_rate: float = 0.0, + attn_drop_rate: float = 0.0, + norm_layer: Callable[..., nn.Module] = partial(LayerNorm, eps=1e-6), + initialize_weights: bool = True, + mask_token: Optional[Parameter] = None, + ): + """Initializes the PixioDecoderTIMM with a deep (default 32-block) decoder.""" + super().__init__( + num_patches=num_patches, + patch_size=patch_size, + in_chans=in_chans, + embed_dim=embed_dim, + decoder_embed_dim=decoder_embed_dim, + decoder_depth=decoder_depth, + decoder_num_heads=decoder_num_heads, + num_prefix_tokens=num_prefix_tokens, + mlp_ratio=mlp_ratio, + proj_drop_rate=proj_drop_rate, + attn_drop_rate=attn_drop_rate, + norm_layer=norm_layer, + initialize_weights=initialize_weights, + mask_token=mask_token, + ) diff --git a/lightly/models/utils.py b/lightly/models/utils.py index eb6c17533..74f739329 100644 --- a/lightly/models/utils.py +++ b/lightly/models/utils.py @@ -612,6 +612,114 @@ def random_token_mask( return idx_keep, idx_mask +def random_grid_token_mask( + size: Tuple[int, int], + mask_ratio: float = 0.75, + grid_size: int = 4, + num_prefix_tokens: int = 1, + device: Optional[Union[torch.device, str]] = None, +) -> Tuple[Tensor, Tensor]: + """Creates random token masks at a coarse grid granularity. + + Instead of masking individual patches (as in :func:`random_token_mask`), whole + ``grid_size`` x ``grid_size`` blocks of patches on a regular grid are kept or + masked together. This is the larger masking granularity used by Pixio [0] to + avoid trivial pixel-reconstruction shortcuts between neighboring patches. + + - [0]: In Pursuit of Pixel Supervision for Visual Pre-training, 2025, + https://arxiv.org/abs/2512.15715 + + Args: + size: + Size of the token batch, (batch_size, sequence_length). The sequence + length must equal num_prefix_tokens + num_patches where num_patches is a + perfect square. + mask_ratio: + Proportion of grid cells to mask. + grid_size: + Side length of a grid cell measured in patches. The patch grid height + (and width) must be divisible by grid_size. + num_prefix_tokens: + Number of prefix tokens (e.g. class or register tokens) that precede the + patch tokens. They are never masked and are always returned first in + idx_keep. + device: + Device on which to create the index masks. + + Returns: + An (idx_keep, idx_mask) tuple of int64 index tensors, where + num_cells = (height // grid_size) ** 2 and + num_keep = int(num_cells * (1 - mask_ratio)) * grid_size ** 2: + + - idx_keep, shape (batch_size, num_prefix_tokens + num_keep): the prefix + token indices followed by the indices of the patches in the kept cells. + - idx_mask, shape (batch_size, num_patches - num_keep): the indices of the + patches in the masked cells. + + Patch index p maps to token index p + num_prefix_tokens. + + Raises: + ValueError: If sequence_length is not greater than num_prefix_tokens, if the + number of patches is not a perfect square, or if the patch grid is not + divisible by grid_size. + """ + batch_size, sequence_length = size + num_patches = sequence_length - num_prefix_tokens + if num_patches <= 0: + raise ValueError( + f"sequence_length ({sequence_length}) must be greater than " + f"num_prefix_tokens ({num_prefix_tokens})." + ) + height = width = int(num_patches**0.5) + if height * width != num_patches: + raise ValueError( + f"Number of patches ({num_patches}) must be a perfect square. Got " + f"sequence_length={sequence_length} and " + f"num_prefix_tokens={num_prefix_tokens}." + ) + if height % grid_size != 0: + raise ValueError( + f"Patch grid side length ({height}) must be divisible by " + f"grid_size ({grid_size})." + ) + + num_cells = (height // grid_size) * (width // grid_size) + num_keep_cells = int(num_cells * (1 - mask_ratio)) + + # Map every grid cell to the patch indices it covers: (1, num_cells, grid_size**2). + patch_indices = torch.arange(num_patches, device=device).reshape(1, height, width) + patch_indices = patch_indices.unfold(1, grid_size, grid_size).unfold( + 2, grid_size, grid_size + ) + patch_indices = patch_indices.reshape(1, num_cells, grid_size * grid_size) + patch_indices = patch_indices.expand(batch_size, -1, -1) + + # Randomly order cells per sample and split into kept and masked cells. + noise = torch.rand(batch_size, num_cells, device=device) # (batch_size, num_cells) + cell_order = torch.argsort(noise, dim=1) # (batch_size, num_cells) + keep_cells = cell_order[:, :num_keep_cells] # (batch_size, num_keep_cells) + mask_cells = cell_order[ + :, num_keep_cells: + ] # (batch_size, num_cells - num_keep_cells) + + # Expand each kept/masked cell to the patch indices it covers, then offset by the + # prefix tokens. keep_patches: (batch_size, num_keep_cells * grid_size**2). + keep_patches = get_at_index(patch_indices, keep_cells).reshape(batch_size, -1) + idx_mask = get_at_index(patch_indices, mask_cells).reshape(batch_size, -1) + keep_patches = keep_patches + num_prefix_tokens + idx_mask = idx_mask + num_prefix_tokens + + # Prefix tokens are always kept and come first: + # idx_keep is (batch_size, num_prefix_tokens + num_keep_cells * grid_size**2). + prefix = torch.arange(num_prefix_tokens, device=device) # (num_prefix_tokens,) + prefix = prefix.unsqueeze(0).expand( + batch_size, -1 + ) # (batch_size, num_prefix_tokens) + idx_keep = torch.cat([prefix, keep_patches], dim=1) + + return idx_keep, idx_mask + + def random_prefix_mask( size: Tuple[int, int], max_prefix_length: int, @@ -1061,7 +1169,7 @@ def initialize_positional_embedding( elif strategy == "sincos": initialize_2d_sine_cosine_positional_embedding( pos_embedding=pos_embedding, - has_class_token=num_prefix_tokens > 0, + num_prefix_tokens=num_prefix_tokens, ) elif strategy == "skip": return @@ -1083,14 +1191,31 @@ def initialize_learnable_positional_embedding(pos_embedding: Parameter) -> None: def initialize_2d_sine_cosine_positional_embedding( - pos_embedding: Parameter, has_class_token: bool + pos_embedding: Parameter, + has_class_token: bool = True, + num_prefix_tokens: Optional[int] = None, ) -> None: + """Initializes a positional embedding in-place with a fixed 2D sine-cosine embedding. + + The positional embedding is frozen (requires_grad set to False) afterwards. + + Args: + pos_embedding: + Positional embedding parameter to initialize in-place. + has_class_token: + Whether the embedding has a single prefix (class) token. Ignored if + num_prefix_tokens is set. + num_prefix_tokens: + Number of prefix tokens (e.g. class or register tokens). Overrides + has_class_token when set. + """ + n_prefix = int(has_class_token) if num_prefix_tokens is None else num_prefix_tokens _, seq_length, hidden_dim = pos_embedding.shape - grid_size = int((seq_length - int(has_class_token)) ** 0.5) + grid_size = int((seq_length - n_prefix) ** 0.5) sine_cosine_embedding = get_2d_sine_cosine_positional_embedding( embed_dim=hidden_dim, grid_size=grid_size, - cls_token=has_class_token, + num_prefix_tokens=n_prefix, ) pos_embedding.data.copy_( torch.from_numpy(sine_cosine_embedding).float().unsqueeze(0) @@ -1100,7 +1225,10 @@ def initialize_2d_sine_cosine_positional_embedding( def get_2d_sine_cosine_positional_embedding( - embed_dim: int, grid_size: int, cls_token: bool + embed_dim: int, + grid_size: int, + cls_token: bool = True, + num_prefix_tokens: Optional[int] = None, ) -> NDArray[np.float32]: """Generates 2D sine-cosine positional embedding. @@ -1112,12 +1240,17 @@ def get_2d_sine_cosine_positional_embedding( grid_size: Height and width of the grid. cls_token: - If True, a positional embedding for the class token is generated. + If True, a single zero positional embedding for the class token is + prepended. Ignored if num_prefix_tokens is set. + num_prefix_tokens: + If set, this many zero positional embeddings are prepended, for models + with multiple class or register tokens. Overrides cls_token. Returns: - Positional embedding with shape (grid_size * grid_size, embed_dim) or - (1 + grid_size * grid_size, embed_dim) if cls_token is True. + Positional embedding with shape (num_prefix + grid_size * grid_size, embed_dim) + where num_prefix is num_prefix_tokens if set, else int(cls_token). """ + n_prefix = int(cls_token) if num_prefix_tokens is None else num_prefix_tokens grid_h = np.arange(grid_size, dtype=np.float32) grid_w = np.arange(grid_size, dtype=np.float32) grid = np.meshgrid(grid_w, grid_h) # here w goes first @@ -1125,8 +1258,8 @@ def get_2d_sine_cosine_positional_embedding( grid = grid.reshape([2, 1, grid_size, grid_size]) pos_embed = get_2d_sine_cosine_positional_embedding_from_grid(embed_dim, grid) - if cls_token: - pos_embed = np.concatenate([np.zeros([1, embed_dim]), pos_embed], axis=0) + if n_prefix > 0: + pos_embed = np.concatenate([np.zeros([n_prefix, embed_dim]), pos_embed], axis=0) return pos_embed diff --git a/tests/models/modules/test_masked_autoencoder_timm.py b/tests/models/modules/test_masked_autoencoder_timm.py index 8d9180963..1f6e45623 100644 --- a/tests/models/modules/test_masked_autoencoder_timm.py +++ b/tests/models/modules/test_masked_autoencoder_timm.py @@ -11,7 +11,7 @@ pytest.skip("TIMM vision transformer is not available", allow_module_level=True) -from lightly.models.modules import MAEDecoderTIMM +from lightly.models.modules import MAEDecoderTIMM, PixioDecoderTIMM class TestMAEDecoderTIMM(unittest.TestCase): @@ -63,3 +63,85 @@ def test_forward(self) -> None: @unittest.skipUnless(torch.cuda.is_available(), "Cuda not available.") def test_forward_cuda(self) -> None: self._test_forward(torch.device("cuda")) + + def test_init__num_prefix_tokens(self) -> None: + num_patches, num_prefix_tokens, decoder_embed_dim = 49, 8, 256 + decoder = MAEDecoderTIMM( + num_patches=num_patches, + patch_size=32, + embed_dim=128, + decoder_embed_dim=decoder_embed_dim, + decoder_depth=2, + decoder_num_heads=4, + num_prefix_tokens=num_prefix_tokens, + ) + self.assertEqual( + list(decoder.decoder_pos_embed.shape), + [1, num_patches + num_prefix_tokens, decoder_embed_dim], + ) + + def _test_forward__num_prefix_tokens(self, device: torch.device) -> None: + torch.manual_seed(0) + num_patches, num_prefix_tokens, embed_input_dim, patch_size = 49, 8, 128, 32 + seq_length = num_patches + num_prefix_tokens + decoder = MAEDecoderTIMM( + num_patches=num_patches, + patch_size=patch_size, + embed_dim=embed_input_dim, + decoder_embed_dim=256, + decoder_depth=2, + decoder_num_heads=4, + num_prefix_tokens=num_prefix_tokens, + ).to(device) + tokens = torch.rand(2, seq_length, embed_input_dim).to(device) + predictions = decoder(tokens) + self.assertListEqual( + list(predictions.shape), [2, seq_length, 3 * patch_size**2] + ) + self.assertTrue(torch.all(torch.not_equal(predictions, torch.inf))) + + def test_forward__num_prefix_tokens(self) -> None: + self._test_forward__num_prefix_tokens(torch.device("cpu")) + + @unittest.skipUnless(torch.cuda.is_available(), "Cuda not available.") + def test_forward__num_prefix_tokens_cuda(self) -> None: + self._test_forward__num_prefix_tokens(torch.device("cuda")) + + +class TestPixioDecoderTIMM: + def test_init__default_depth_is_32(self) -> None: + decoder = PixioDecoderTIMM( + num_patches=256, + patch_size=16, + embed_dim=768, + decoder_embed_dim=512, + decoder_num_heads=16, + num_prefix_tokens=8, + ) + assert len(decoder.decoder_blocks) == 32 + assert list(decoder.decoder_pos_embed.shape) == [1, 256 + 8, 512] + + def _test_forward(self, device: torch.device) -> None: + torch.manual_seed(0) + num_patches, num_prefix_tokens, patch_size = 64, 8, 16 + seq_length = num_patches + num_prefix_tokens + decoder = PixioDecoderTIMM( + num_patches=num_patches, + patch_size=patch_size, + embed_dim=128, + decoder_embed_dim=64, + decoder_depth=2, # keep the test cheap + decoder_num_heads=4, + num_prefix_tokens=num_prefix_tokens, + ).to(device) + tokens = torch.rand(2, seq_length, 128).to(device) + predictions = decoder(tokens) + assert list(predictions.shape) == [2, seq_length, 3 * patch_size**2] + assert torch.all(torch.not_equal(predictions, torch.inf)) + + def test_forward(self) -> None: + self._test_forward(torch.device("cpu")) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="Cuda not available.") + def test_forward_cuda(self) -> None: + self._test_forward(torch.device("cuda")) diff --git a/tests/models/test_ModelUtils.py b/tests/models/test_ModelUtils.py index 8189b36b1..36ff08045 100644 --- a/tests/models/test_ModelUtils.py +++ b/tests/models/test_ModelUtils.py @@ -901,6 +901,20 @@ def test_initialize_positional_embedding( mock_fn.assert_called_once() +def test_initialize_positional_embedding__sincos_multiple_prefix_tokens() -> None: + num_prefix_tokens, grid_size, embed_dim = 8, 4, 16 + seq_length = num_prefix_tokens + grid_size * grid_size + pos_embedding = Parameter(torch.rand(1, seq_length, embed_dim)) + utils.initialize_positional_embedding( + pos_embedding=pos_embedding, + strategy="sincos", + num_prefix_tokens=num_prefix_tokens, + ) + # prefix rows are zeroed by the sincos init + assert (pos_embedding[0, :num_prefix_tokens] == 0).all() + assert not (pos_embedding[0, num_prefix_tokens:] == 0).all() + + def test_initialize_learnable_positional_embedding() -> None: pos_embedding = Parameter(torch.ones(1, 1, 64)) orig_pos_embedding = pos_embedding.clone() @@ -918,10 +932,13 @@ def test_normalize_mean_var() -> None: assert norm[1] == pytest.approx(0.0) assert norm[2] == pytest.approx(1) + # seed for determinism; atol is loosened because normalize_mean_var regularizes + # with eps, so the output variance is 1 - eps / var(x) rather than exactly 1. + torch.manual_seed(0) x = torch.rand(2, 3, 4) norm = utils.normalize_mean_var(x) assert torch.allclose(norm.mean(dim=-1), torch.tensor(0.0), rtol=0.0001, atol=1e-5) - assert torch.allclose(norm.var(dim=-1), torch.tensor(1.0), rtol=0.0001, atol=1e-5) + assert torch.allclose(norm.var(dim=-1), torch.tensor(1.0), rtol=0.0001, atol=1e-4) def test_update_drop_path_rate__uniform() -> None: @@ -971,3 +988,168 @@ def test_update_drop_path_rate__unknown_mode() -> None: model = VisionTransformer(drop_path_rate=0, depth=4) with pytest.raises(ValueError, match="Unknown mode"): utils.update_drop_path_rate(model=model, drop_path_rate=0.1, mode="unknown") + + +def test_random_grid_token_mask__partition() -> None: + torch.manual_seed(0) + batch_size, num_prefix_tokens = 2, 1 + sequence_length = num_prefix_tokens + 16 # 4x4 patch grid + idx_keep, idx_mask = utils.random_grid_token_mask( + size=(batch_size, sequence_length), + mask_ratio=0.5, + grid_size=2, + num_prefix_tokens=num_prefix_tokens, + ) + idx, _ = torch.cat([idx_keep, idx_mask], dim=1).sort(dim=1) + expected = torch.arange(sequence_length).expand(batch_size, sequence_length) + assert torch.equal(idx, expected) + + +def test_random_grid_token_mask__prefix_tokens_kept() -> None: + torch.manual_seed(0) + batch_size, num_prefix_tokens = 3, 8 + sequence_length = num_prefix_tokens + 16 + idx_keep, idx_mask = utils.random_grid_token_mask( + size=(batch_size, sequence_length), + mask_ratio=0.75, + grid_size=2, + num_prefix_tokens=num_prefix_tokens, + ) + prefix = torch.arange(num_prefix_tokens).expand(batch_size, num_prefix_tokens) + assert torch.equal(idx_keep[:, :num_prefix_tokens], prefix) + assert idx_mask.min().item() >= num_prefix_tokens + + +def test_random_grid_token_mask__whole_cell_granularity() -> None: + torch.manual_seed(0) + num_prefix_tokens = 1 + sequence_length = num_prefix_tokens + 16 + cells = [{0, 1, 4, 5}, {2, 3, 6, 7}, {8, 9, 12, 13}, {10, 11, 14, 15}] + idx_keep, idx_mask = utils.random_grid_token_mask( + size=(1, sequence_length), + mask_ratio=0.5, + grid_size=2, + num_prefix_tokens=num_prefix_tokens, + ) + masked = {int(i) - num_prefix_tokens for i in idx_mask[0]} + for cell in cells: + assert masked.issuperset(cell) or masked.isdisjoint(cell) + + +def test_random_grid_token_mask__keep_count() -> None: + num_prefix_tokens = 1 + sequence_length = num_prefix_tokens + 16 # 4 cells of 4 patches (grid=2) + idx_keep, idx_mask = utils.random_grid_token_mask( + size=(2, sequence_length), + mask_ratio=0.75, + grid_size=2, + num_prefix_tokens=num_prefix_tokens, + ) + # keep int(4 * 0.25) = 1 cell -> 4 patches, + 1 prefix + assert idx_keep.shape == (2, num_prefix_tokens + 4) + assert idx_mask.shape == (2, 12) + + +def test_random_grid_token_mask__not_divisible_raises() -> None: + with pytest.raises(ValueError): + utils.random_grid_token_mask( + size=(1, 1 + 16), mask_ratio=0.5, grid_size=3, num_prefix_tokens=1 + ) + + +def test_random_grid_token_mask__not_square_raises() -> None: + with pytest.raises(ValueError): + utils.random_grid_token_mask( + size=(1, 1 + 15), mask_ratio=0.5, grid_size=2, num_prefix_tokens=1 + ) + + +def test_random_grid_token_mask__prefix_ge_sequence_raises() -> None: + with pytest.raises(ValueError): + # num_prefix_tokens >= sequence_length -> non-positive patch count. + utils.random_grid_token_mask( + size=(1, 3), mask_ratio=0.5, grid_size=2, num_prefix_tokens=5 + ) + + +def test_random_grid_token_mask__mask_ratio_extremes() -> None: + num_prefix_tokens = 1 + sequence_length = num_prefix_tokens + 16 # 4 cells of 4 patches (grid=2) + # mask_ratio=0.0 keeps all patches and masks none. + idx_keep, idx_mask = utils.random_grid_token_mask( + size=(2, sequence_length), + mask_ratio=0.0, + grid_size=2, + num_prefix_tokens=num_prefix_tokens, + ) + assert idx_keep.shape == (2, sequence_length) + assert idx_mask.shape == (2, 0) + # mask_ratio=1.0 keeps only the prefix tokens and masks all patches. + idx_keep, idx_mask = utils.random_grid_token_mask( + size=(2, sequence_length), + mask_ratio=1.0, + grid_size=2, + num_prefix_tokens=num_prefix_tokens, + ) + assert idx_keep.shape == (2, num_prefix_tokens) + assert idx_mask.shape == (2, 16) + + +def test_random_grid_token_mask__grid_size_4() -> None: + torch.manual_seed(0) + num_prefix_tokens = 8 + # 16x16 patch grid -> 4x4 = 16 cells of 4x4 patches (Pixio's headline config). + sequence_length = num_prefix_tokens + 256 + idx_keep, idx_mask = utils.random_grid_token_mask( + size=(2, sequence_length), + mask_ratio=0.75, + grid_size=4, + num_prefix_tokens=num_prefix_tokens, + ) + # keep int(16 * 0.25) = 4 cells -> 64 patches (+ 8 prefix); mask 12 cells -> 192. + assert idx_keep.shape == (2, num_prefix_tokens + 64) + assert idx_mask.shape == (2, 192) + idx, _ = torch.cat([idx_keep, idx_mask], dim=1).sort(dim=1) + expected = torch.arange(sequence_length).expand(2, sequence_length) + assert torch.equal(idx, expected) + + +@pytest.mark.parametrize("device", ["cpu", "cuda"]) # type: ignore[misc] +def test_random_grid_token_mask__device(device: str) -> None: + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + + idx_keep, idx_mask = utils.random_grid_token_mask( + size=(2, 1 + 16), + mask_ratio=0.5, + grid_size=2, + num_prefix_tokens=1, + device=device, + ) + assert idx_keep.device.type == device + assert idx_mask.device.type == device + + +def test_get_2d_sine_cosine_positional_embedding__num_prefix_tokens() -> None: + embed_dim, grid_size = 16, 4 + emb = utils.get_2d_sine_cosine_positional_embedding( + embed_dim=embed_dim, grid_size=grid_size, num_prefix_tokens=8 + ) + # 8 prefix rows + grid_size**2 patch rows + assert emb.shape == (8 + grid_size * grid_size, embed_dim) + # prefix rows are zeros + assert (emb[:8] == 0).all() + # patch rows are not all zero + assert not (emb[8:] == 0).all() + + +def test_get_2d_sine_cosine_positional_embedding__backwards_compatible() -> None: + embed_dim, grid_size = 16, 4 + with_cls = utils.get_2d_sine_cosine_positional_embedding( + embed_dim=embed_dim, grid_size=grid_size, cls_token=True + ) + one_prefix = utils.get_2d_sine_cosine_positional_embedding( + embed_dim=embed_dim, grid_size=grid_size, num_prefix_tokens=1 + ) + assert with_cls.shape == one_prefix.shape == (1 + grid_size * grid_size, embed_dim) + assert (with_cls == one_prefix).all()