diff --git a/README.md b/README.md index e1412f625..c2a31c455 100644 --- a/README.md +++ b/README.md @@ -77,6 +77,7 @@ and PyTorch Lightning distributed examples for all models to kickstart your proj | SimSiam | 2021 | [paper](https://arxiv.org/abs/2011.10566) | [docs](https://docs.lightly.ai/self-supervised-learning/examples/simsiam.html) | [![Open In Colab](https://img.shields.io/badge/Colab-PyTorch-blue?logo=googlecolab)](https://colab.research.google.com/github/lightly-ai/lightly/blob/master/examples/notebooks/pytorch/simsiam.ipynb) | [![Open In Colab](https://img.shields.io/badge/Colab-PyTorch_Lightning-blue?logo=googlecolab)](https://colab.research.google.com/github/lightly-ai/lightly/blob/master/examples/notebooks/pytorch_lightning/simsiam.ipynb) | | SwaV | 2020 | [paper](https://arxiv.org/abs/2006.09882) | [docs](https://docs.lightly.ai/self-supervised-learning/examples/swav.html) | [![Open In Colab](https://img.shields.io/badge/Colab-PyTorch-blue?logo=googlecolab)](https://colab.research.google.com/github/lightly-ai/lightly/blob/master/examples/notebooks/pytorch/swav.ipynb) | [![Open In Colab](https://img.shields.io/badge/Colab-PyTorch_Lightning-blue?logo=googlecolab)](https://colab.research.google.com/github/lightly-ai/lightly/blob/master/examples/notebooks/pytorch_lightning/swav.ipynb) | | VICReg | 2021 | [paper](https://arxiv.org/abs/2105.04906) | [docs](https://docs.lightly.ai/self-supervised-learning/examples/vicreg.html) | [![Open In Colab](https://img.shields.io/badge/Colab-PyTorch-blue?logo=googlecolab)](https://colab.research.google.com/github/lightly-ai/lightly/blob/master/examples/notebooks/pytorch/vicreg.ipynb) | [![Open In Colab](https://img.shields.io/badge/Colab-PyTorch_Lightning-blue?logo=googlecolab)](https://colab.research.google.com/github/lightly-ai/lightly/blob/master/examples/notebooks/pytorch_lightning/vicreg.ipynb) | +| BEiT | 2021 | [paper](https://arxiv.org/pdf/2106.08254) | [docs](https://docs.lightly.ai/self-supervised-learning/examples/beit.html) | [![Open In Colab](https://img.shields.io/badge/Colab-PyTorch-blue?logo=googlecolab)](https://colab.research.google.com/github/lightly-ai/lightly/blob/master/examples/notebooks/pytorch/beit.ipynb) | [![Open In Colab](https://img.shields.io/badge/Colab-PyTorch_Lightning-blue?logo=googlecolab)](https://colab.research.google.com/github/lightly-ai/lightly/blob/master/examples/notebooks/pytorch_lightning/beit.ipynb) | ## Tutorials diff --git a/docs/source/examples/beit.rst b/docs/source/examples/beit.rst new file mode 100644 index 000000000..29bbab873 --- /dev/null +++ b/docs/source/examples/beit.rst @@ -0,0 +1,74 @@ +.. _beit: + +BEiT +==== + +Implementation of the BEiT (BERT Pre-Training of Image Transformers) +architecture for masked image modeling (MIM). + +BEiT pre-trains a Vision Transformer by masking random patches of the input +image and predicting the discrete visual tokens of the masked patches. Unlike +MAE which predicts raw pixels, BEiT uses a pre-trained discrete VAE tokenizer +to convert images into a vocabulary of visual tokens, and the transformer learns +to predict these token indices. + +Key components: + +- **Blockwise masking**: Random block-shaped regions are masked rather than + individual patches, following the BEIT paper. +- **Discrete VAE tokenizer**: A pre-trained tokenizer (e.g., DALL-E dVAE) + converts images into visual token targets. +- **Learnable mask token**: A special [M] embedding replaces masked patches + in the input. +- **Shared relative position bias**: Optional relative position bias can be + used instead of absolute position embeddings. + +Reference: + `BEIT: BERT Pre-Training of Image Transformers, 2021 `_ + + +.. 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/beit.ipynb + + This example can be run from the command line with:: + + python lightly/examples/pytorch/beit.py + + .. literalinclude:: ../../../examples/pytorch/beit.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/beit.ipynb + + This example can be run from the command line with:: + + python lightly/examples/pytorch_lightning/beit.py + + .. literalinclude:: ../../../examples/pytorch_lightning/beit.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/beit.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/beit.py + + The model differs in the following ways from the non-distributed + implementation: + + - Distributed Data Parallel is enabled + - Synchronized Batch Norm is used in place of standard Batch Norm + + Note that Synchronized Batch Norm is optional and the model can also be + trained without it. Without Synchronized Batch Norm the batch norm for + each GPU is only calculated based on the features on that specific GPU. + + .. literalinclude:: ../../../examples/pytorch_lightning_distributed/beit.py \ No newline at end of file diff --git a/docs/source/examples/models.rst b/docs/source/examples/models.rst index b6d24b02e..ef606c88b 100644 --- a/docs/source/examples/models.rst +++ b/docs/source/examples/models.rst @@ -35,3 +35,4 @@ for PyTorch and PyTorch Lightning to give you a headstart when implementing your tico.rst vicreg.rst vicregl.rst + beit.rst diff --git a/examples/notebooks/pytorch/beit.ipynb b/examples/notebooks/pytorch/beit.ipynb new file mode 100644 index 000000000..608c61045 --- /dev/null +++ b/examples/notebooks/pytorch/beit.ipynb @@ -0,0 +1,273 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "0", + "metadata": { + "lines_to_next_cell": 2 + }, + "source": [ + "Example: BEIT pre-training on CIFAR-10.\n", + "This script demonstrates how to use the BEIT model for masked image\n", + "modeling (MIM) pre-training on a small subset of CIFAR-10.\n", + "NOTE:\n", + "The tokenizer must be pretrained before running BEiT pretraining.\n", + "A pretrained tokenizer should be loaded from a checkpoint." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "1", + "metadata": {}, + "outputs": [], + "source": [ + "from __future__ import annotations" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "2", + "metadata": {}, + "outputs": [], + "source": [ + "import torch\n", + "import torchvision\n", + "from torch import nn\n", + "from torch.utils.data import Subset" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "3", + "metadata": {}, + "outputs": [], + "source": [ + "from lightly.loss import MaskedImageModelingLoss\n", + "from lightly.models.modules import BEITEncoder, BEiTImageTokenizer, BEiTMIMHead\n", + "from lightly.transforms import BEiTTransform" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4", + "metadata": {}, + "outputs": [], + "source": [ + "class BEiT(nn.Module):\n", + " \"\"\"Simple BEIT pre-training wrapper.\n", + "\n", + " Attributes:\n", + " encoder:\n", + " BEIT Vision Transformer encoder.\n", + " tokenizer:\n", + " Discrete visual tokenizer (e.g., dVAE or VQ-VAE).\n", + " head:\n", + " Linear projection head mapping patch features to vocabulary\n", + " logits.\n", + " \"\"\"\n", + "\n", + " def __init__(\n", + " self,\n", + " encoder: BEITEncoder,\n", + " tokenizer: BEiTImageTokenizer,\n", + " ) -> None:\n", + " \"\"\"Initializes the BEIT model.\n", + "\n", + " Args:\n", + " encoder:\n", + " BEIT encoder instance.\n", + " tokenizer:\n", + " Visual tokenizer for generating target token IDs.\n", + " \"\"\"\n", + " super().__init__()\n", + " self.encoder = encoder\n", + " self.tokenizer = tokenizer\n", + " self.head = BEiTMIMHead(\n", + " embed_dim=encoder.embed_dim,\n", + " vocab_size=tokenizer.vocab_size,\n", + " )\n", + "\n", + " def forward(\n", + " self,\n", + " images: torch.Tensor,\n", + " bool_masked_pos: torch.BoolTensor,\n", + " ) -> tuple[torch.Tensor, torch.Tensor]:\n", + " \"\"\"Forward pass for masked image modeling.\n", + "\n", + " Args:\n", + " images:\n", + " Input images of shape (B, C, H, W).\n", + " bool_masked_pos:\n", + " Boolean mask of shape (B, N) indicating masked patches.\n", + "\n", + " Returns:\n", + " A tuple of:\n", + " - mim_logits: Logits for masked positions.\n", + " - token_targets: Target token IDs for masked positions.\n", + " \"\"\"\n", + " enc_out = self.encoder(\n", + " x=images,\n", + " bool_masked_pos=bool_masked_pos,\n", + " )\n", + " patch_features = enc_out[\"patch_features\"]\n", + "\n", + " all_logits = self.head(patch_features=patch_features)\n", + " mim_logits = all_logits[bool_masked_pos]\n", + "\n", + " with torch.no_grad():\n", + " token_ids = self.tokenizer.tokenize(x=images)\n", + " token_targets = token_ids[bool_masked_pos]\n", + "\n", + " return mim_logits, token_targets" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "5", + "metadata": {}, + "outputs": [], + "source": [ + "# Smaller encoder for fast experimentation: depth=2, embed_dim=256\n", + "encoder = BEITEncoder(\n", + " img_size=224,\n", + " patch_size=8,\n", + " embed_dim=256,\n", + " depth=2,\n", + " num_heads=4,\n", + ")\n", + "tokenizer = BEiTImageTokenizer(vocab_size=512)\n", + "# disables gradients for every parameter in the tokenizer\n", + "tokenizer.requires_grad_(False)\n", + "tokenizer.eval()\n", + "model = BEiT(encoder=encoder, tokenizer=tokenizer)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "6", + "metadata": {}, + "outputs": [], + "source": [ + "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n", + "model.to(device=device)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "7", + "metadata": {}, + "outputs": [], + "source": [ + "transform = BEiTTransform(input_size=224, patch_size=8)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "8", + "metadata": {}, + "outputs": [], + "source": [ + "# CIFAR-10: downloads fast, small images\n", + "dataset = torchvision.datasets.CIFAR10(\n", + " root=\"datasets/cifar10\",\n", + " train=True,\n", + " download=True,\n", + " transform=transform,\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "9", + "metadata": {}, + "outputs": [], + "source": [ + "# Only use 200 samples for quick demonstration\n", + "dataset = Subset(dataset, indices=list(range(200)))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "10", + "metadata": {}, + "outputs": [], + "source": [ + "dataloader = torch.utils.data.DataLoader(\n", + " dataset=dataset,\n", + " batch_size=16,\n", + " shuffle=True,\n", + " drop_last=True,\n", + " num_workers=2,\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "11", + "metadata": {}, + "outputs": [], + "source": [ + "criterion = MaskedImageModelingLoss()\n", + "optimizer = torch.optim.AdamW(\n", + " params=model.parameters(),\n", + " lr=1.5e-3,\n", + " weight_decay=0.05,\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "12", + "metadata": {}, + "outputs": [], + "source": [ + "print(\"Starting Training\")\n", + "model.train()\n", + "for epoch in range(2):\n", + " total_loss = 0.0\n", + " for batch in dataloader:\n", + " images, _ = batch\n", + " images = images.to(device=device)\n", + "\n", + " bool_masked_pos = transform.mask_generator(\n", + " batch_size=images.shape[0],\n", + " ).to(device=device)\n", + "\n", + " mim_logits, token_targets = model(\n", + " images=images,\n", + " bool_masked_pos=bool_masked_pos,\n", + " )\n", + " loss = criterion(mim_logits, token_targets)\n", + " total_loss += loss.item()\n", + "\n", + " loss.backward()\n", + " optimizer.step()\n", + " optimizer.zero_grad()\n", + "\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/beit.ipynb b/examples/notebooks/pytorch_lightning/beit.ipynb new file mode 100644 index 000000000..e8019581d --- /dev/null +++ b/examples/notebooks/pytorch_lightning/beit.ipynb @@ -0,0 +1,275 @@ +{ + "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": "markdown", + "id": "2", + "metadata": {}, + "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." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "3", + "metadata": {}, + "outputs": [], + "source": [ + "import pytorch_lightning as pl\n", + "import torch\n", + "import torchvision\n", + "from torch import nn\n", + "from torch.utils.data import Subset" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4", + "metadata": {}, + "outputs": [], + "source": [ + "from lightly.loss import MaskedImageModelingLoss\n", + "from lightly.models.modules import BEITEncoder, BEiTImageTokenizer, BEiTMIMHead\n", + "from lightly.transforms import BEiTTransform" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "5", + "metadata": {}, + "outputs": [], + "source": [ + "class BEIT(pl.LightningModule):\n", + " \"\"\"BEIT pre-training module for masked image modeling.\n", + "\n", + " Attributes:\n", + " backbone:\n", + " BEIT Vision Transformer encoder.\n", + " projection_head:\n", + " MIM head mapping patch features to vocabulary logits.\n", + " tokenizer:\n", + " Discrete visual tokenizer for generating target token IDs.\n", + " criterion:\n", + " Cross-entropy loss for masked token prediction.\n", + " \"\"\"\n", + "\n", + " def __init__(\n", + " self,\n", + " img_size: int = 224,\n", + " patch_size: int = 8,\n", + " embed_dim: int = 256,\n", + " depth: int = 2,\n", + " num_heads: int = 4,\n", + " vocab_size: int = 512,\n", + " ) -> None:\n", + " \"\"\"Initializes the BEIT pre-training module.\n", + "\n", + " Args:\n", + " img_size:\n", + " Spatial resolution of input images.\n", + " patch_size:\n", + " Size of each patch.\n", + " embed_dim:\n", + " Dimension of token embeddings.\n", + " depth:\n", + " Number of Transformer blocks.\n", + " num_heads:\n", + " Number of attention heads per block.\n", + " vocab_size:\n", + " Size of the discrete visual vocabulary.\n", + " \"\"\"\n", + " super().__init__()\n", + " self.backbone = BEITEncoder(\n", + " img_size=img_size,\n", + " patch_size=patch_size,\n", + " embed_dim=embed_dim,\n", + " depth=depth,\n", + " num_heads=num_heads,\n", + " )\n", + " self.num_patches = self.backbone.num_patches\n", + " self.projection_head = BEiTMIMHead(\n", + " embed_dim=embed_dim,\n", + " vocab_size=vocab_size,\n", + " )\n", + " self.tokenizer = BEiTImageTokenizer(vocab_size=vocab_size)\n", + " self.tokenizer.requires_grad_(False)\n", + " self.tokenizer.eval()\n", + " self.criterion = MaskedImageModelingLoss()\n", + "\n", + " def training_step(\n", + " self,\n", + " batch: tuple[torch.Tensor, torch.Tensor],\n", + " batch_idx: int,\n", + " ) -> torch.Tensor:\n", + " \"\"\"Performs a single training step.\n", + "\n", + " Args:\n", + " batch:\n", + " Tuple of (images, targets) where images is a tensor of\n", + " shape (B, C, H, W).\n", + " batch_idx:\n", + " Index of the current batch.\n", + "\n", + " Returns:\n", + " The computed loss value.\n", + " \"\"\"\n", + " images, _ = batch\n", + " batch_size = images.shape[0]\n", + "\n", + " # Generate blockwise mask\n", + " bool_masked_pos = self.transform.mask_generator(\n", + " batch_size=batch_size,\n", + " ).to(device=images.device)\n", + "\n", + " # Forward pass through encoder + head\n", + " enc_out = self.backbone(\n", + " x=images,\n", + " bool_masked_pos=bool_masked_pos,\n", + " )\n", + " patch_features = enc_out[\"patch_features\"]\n", + " all_logits = self.projection_head(patch_features=patch_features)\n", + "\n", + " # Select masked positions\n", + " mim_logits = all_logits[bool_masked_pos]\n", + "\n", + " # Get token targets from frozen tokenizer\n", + " with torch.no_grad():\n", + " token_ids = self.tokenizer.tokenize(x=images)\n", + " token_targets = token_ids[bool_masked_pos]\n", + "\n", + " loss = self.criterion(mim_logits, token_targets)\n", + " return loss\n", + "\n", + " def configure_optimizers(self) -> torch.optim.Optimizer:\n", + " \"\"\"Configures the optimizer for training.\n", + "\n", + " Returns:\n", + " AdamW optimizer.\n", + " \"\"\"\n", + " optim = torch.optim.AdamW(\n", + " params=self.parameters(),\n", + " lr=1.5e-3,\n", + " weight_decay=0.05,\n", + " )\n", + " return optim" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "6", + "metadata": {}, + "outputs": [], + "source": [ + "model = BEIT()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "7", + "metadata": {}, + "outputs": [], + "source": [ + "transform = BEiTTransform(input_size=224, patch_size=8)\n", + "model.transform = transform # Store transform for access in training_step" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "8", + "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": "9", + "metadata": {}, + "outputs": [], + "source": [ + "dataset = torchvision.datasets.FakeData(\n", + " size=200, # number of fake samples\n", + " image_size=(3, 224, 224),\n", + " num_classes=10,\n", + " transform=transform,\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "10", + "metadata": {}, + "outputs": [], + "source": [ + "dataloader = torch.utils.data.DataLoader(\n", + " dataset=dataset,\n", + " batch_size=16,\n", + " shuffle=True,\n", + " drop_last=True,\n", + " num_workers=2,\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "11", + "metadata": {}, + "outputs": [], + "source": [ + "accelerator = \"gpu\" if torch.cuda.is_available() else \"cpu\"" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "12", + "metadata": {}, + "outputs": [], + "source": [ + "trainer = pl.Trainer(max_epochs=10, devices=1, accelerator=accelerator)\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/beit.ipynb b/examples/notebooks/pytorch_lightning_distributed/beit.ipynb new file mode 100644 index 000000000..390a4b42f --- /dev/null +++ b/examples/notebooks/pytorch_lightning_distributed/beit.ipynb @@ -0,0 +1,245 @@ +{ + "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": "markdown", + "id": "2", + "metadata": {}, + "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." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "3", + "metadata": {}, + "outputs": [], + "source": [ + "import pytorch_lightning as pl\n", + "import torch\n", + "import torchvision" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4", + "metadata": {}, + "outputs": [], + "source": [ + "from lightly.loss import MaskedImageModelingLoss\n", + "from lightly.models.modules import BEITEncoder, BEiTImageTokenizer, BEiTMIMHead\n", + "from lightly.transforms import BEiTTransform" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "5", + "metadata": {}, + "outputs": [], + "source": [ + "class BEiT(pl.LightningModule):\n", + " \"\"\"BEiT pre-training module for masked image modeling with DDP support.\n", + "\n", + " Attributes:\n", + " backbone:\n", + " BEiT Vision Transformer encoder.\n", + " projection_head:\n", + " MIM head mapping patch features to vocabulary logits.\n", + " tokenizer:\n", + " Discrete visual tokenizer for generating target token IDs.\n", + " criterion:\n", + " Cross-entropy loss for masked token prediction.\n", + " \"\"\"\n", + "\n", + " def __init__(\n", + " self,\n", + " img_size: int = 224,\n", + " patch_size: int = 8,\n", + " embed_dim: int = 256,\n", + " depth: int = 2,\n", + " num_heads: int = 4,\n", + " vocab_size: int = 512,\n", + " ) -> None:\n", + " super().__init__()\n", + " self.backbone = BEITEncoder(\n", + " img_size=img_size,\n", + " patch_size=patch_size,\n", + " embed_dim=embed_dim,\n", + " depth=depth,\n", + " num_heads=num_heads,\n", + " )\n", + " self.num_patches = self.backbone.num_patches\n", + " self.projection_head = BEiTMIMHead(\n", + " embed_dim=embed_dim,\n", + " vocab_size=vocab_size,\n", + " )\n", + " self.tokenizer = BEiTImageTokenizer(vocab_size=vocab_size)\n", + " self.tokenizer.requires_grad_(False)\n", + " self.tokenizer.eval()\n", + " self.criterion = MaskedImageModelingLoss()\n", + "\n", + " def training_step(\n", + " self,\n", + " batch: tuple[torch.Tensor, torch.Tensor],\n", + " batch_idx: int,\n", + " ) -> torch.Tensor:\n", + " images, _ = batch\n", + " batch_size = images.shape[0]\n", + "\n", + " bool_masked_pos = self.transform.mask_generator(\n", + " batch_size=batch_size,\n", + " ).to(device=images.device)\n", + "\n", + " enc_out = self.backbone(\n", + " x=images,\n", + " bool_masked_pos=bool_masked_pos,\n", + " )\n", + " patch_features = enc_out[\"patch_features\"]\n", + " all_logits = self.projection_head(patch_features=patch_features)\n", + "\n", + " mim_logits = all_logits[bool_masked_pos]\n", + "\n", + " with torch.no_grad():\n", + " token_ids = self.tokenizer.tokenize(x=images)\n", + " token_targets = token_ids[bool_masked_pos]\n", + "\n", + " loss = self.criterion(mim_logits, token_targets)\n", + " return loss\n", + "\n", + " def configure_optimizers(self) -> torch.optim.Optimizer:\n", + " optim = torch.optim.AdamW(\n", + " params=self.parameters(),\n", + " lr=1.5e-3,\n", + " weight_decay=0.05,\n", + " )\n", + " return optim" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "6", + "metadata": {}, + "outputs": [], + "source": [ + "model = BEiT()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "7", + "metadata": {}, + "outputs": [], + "source": [ + "transform = BEiTTransform(input_size=224, patch_size=8)\n", + "model.transform = transform # Store transform for access in training_step" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "8", + "metadata": {}, + "outputs": [], + "source": [ + "# Fast\n", + "dataset = torchvision.datasets.FakeData(\n", + " size=200,\n", + " image_size=(3, 224, 224),\n", + " num_classes=20,\n", + " transform=transform,\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "9", + "metadata": {}, + "outputs": [], + "source": [ + "dataloader = torch.utils.data.DataLoader(\n", + " dataset,\n", + " batch_size=16,\n", + " shuffle=True,\n", + " drop_last=True,\n", + " num_workers=2,\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "10", + "metadata": { + "lines_to_next_cell": 2 + }, + "outputs": [], + "source": [ + "accelerator = \"gpu\" if torch.cuda.is_available() else \"cpu\"\n", + "n_gpus = torch.cuda.device_count()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "11", + "metadata": { + "lines_to_next_cell": 2 + }, + "outputs": [], + "source": [ + "trainer = pl.Trainer(\n", + " max_epochs=10,\n", + " devices=\"auto\",\n", + " accelerator=\"gpu\",\n", + " sync_batchnorm=True,\n", + " use_distributed_sampler=True, # or replace_sampler_ddp=True for PyTorch Lightning <2.0\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "12", + "metadata": {}, + "outputs": [], + "source": [ + "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/beit.py b/examples/pytorch/beit.py new file mode 100644 index 000000000..bfb0ac304 --- /dev/null +++ b/examples/pytorch/beit.py @@ -0,0 +1,158 @@ +# Example: BEIT pre-training on CIFAR-10. +# This script demonstrates how to use the BEIT model for masked image +# modeling (MIM) pre-training on a small subset of CIFAR-10. +# NOTE: +# The tokenizer must be pretrained before running BEiT pretraining. +# A pretrained tokenizer should be loaded from a checkpoint. + + +from __future__ import annotations + +import torch +import torchvision +from torch import nn +from torch.utils.data import Subset + +from lightly.loss import MaskedImageModelingLoss +from lightly.models.modules import BEITEncoder, BEiTImageTokenizer, BEiTMIMHead +from lightly.transforms import BEiTTransform + + +class BEiT(nn.Module): + """Simple BEIT pre-training wrapper. + + Attributes: + encoder: + BEIT Vision Transformer encoder. + tokenizer: + Discrete visual tokenizer (e.g., dVAE or VQ-VAE). + head: + Linear projection head mapping patch features to vocabulary + logits. + """ + + def __init__( + self, + encoder: BEITEncoder, + tokenizer: BEiTImageTokenizer, + ) -> None: + """Initializes the BEIT model. + + Args: + encoder: + BEIT encoder instance. + tokenizer: + Visual tokenizer for generating target token IDs. + """ + super().__init__() + self.encoder = encoder + self.tokenizer = tokenizer + self.head = BEiTMIMHead( + embed_dim=encoder.embed_dim, + vocab_size=tokenizer.vocab_size, + ) + + def forward( + self, + images: torch.Tensor, + bool_masked_pos: torch.BoolTensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Forward pass for masked image modeling. + + Args: + images: + Input images of shape (B, C, H, W). + bool_masked_pos: + Boolean mask of shape (B, N) indicating masked patches. + + Returns: + A tuple of: + - mim_logits: Logits for masked positions. + - token_targets: Target token IDs for masked positions. + """ + enc_out = self.encoder( + x=images, + bool_masked_pos=bool_masked_pos, + ) + patch_features = enc_out["patch_features"] + + all_logits = self.head(patch_features=patch_features) + mim_logits = all_logits[bool_masked_pos] + + with torch.no_grad(): + token_ids = self.tokenizer.tokenize(x=images) + token_targets = token_ids[bool_masked_pos] + + return mim_logits, token_targets + + +# Smaller encoder for fast experimentation: depth=2, embed_dim=256 +encoder = BEITEncoder( + img_size=224, + patch_size=8, + embed_dim=256, + depth=2, + num_heads=4, +) +tokenizer = BEiTImageTokenizer(vocab_size=512) +# disables gradients for every parameter in the tokenizer +tokenizer.requires_grad_(False) +tokenizer.eval() +model = BEiT(encoder=encoder, tokenizer=tokenizer) + +device = "cuda" if torch.cuda.is_available() else "cpu" +model.to(device=device) + +transform = BEiTTransform(input_size=224, patch_size=8) + +# CIFAR-10: downloads fast, small images +dataset = torchvision.datasets.CIFAR10( + root="datasets/cifar10", + train=True, + download=True, + transform=transform, +) + +# Only use 200 samples for quick demonstration +dataset = Subset(dataset, indices=list(range(200))) + +dataloader = torch.utils.data.DataLoader( + dataset=dataset, + batch_size=16, + shuffle=True, + drop_last=True, + num_workers=2, +) + +criterion = MaskedImageModelingLoss() +optimizer = torch.optim.AdamW( + params=model.parameters(), + lr=1.5e-3, + weight_decay=0.05, +) + +print("Starting Training") +model.train() +for epoch in range(2): + total_loss = 0.0 + for batch in dataloader: + images, _ = batch + images = images.to(device=device) + + bool_masked_pos = transform.mask_generator( + batch_size=images.shape[0], + ).to(device=device) + + mim_logits, token_targets = model( + images=images, + bool_masked_pos=bool_masked_pos, + ) + loss = criterion(mim_logits, token_targets) + total_loss += loss.item() + + 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/beit.py b/examples/pytorch_lightning/beit.py new file mode 100644 index 000000000..ae28c7f86 --- /dev/null +++ b/examples/pytorch_lightning/beit.py @@ -0,0 +1,163 @@ +# 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 torch import nn +from torch.utils.data import Subset + +from lightly.loss import MaskedImageModelingLoss +from lightly.models.modules import BEITEncoder, BEiTImageTokenizer, BEiTMIMHead +from lightly.transforms import BEiTTransform + + +class BEIT(pl.LightningModule): + """BEIT pre-training module for masked image modeling. + + Attributes: + backbone: + BEIT Vision Transformer encoder. + projection_head: + MIM head mapping patch features to vocabulary logits. + tokenizer: + Discrete visual tokenizer for generating target token IDs. + criterion: + Cross-entropy loss for masked token prediction. + """ + + def __init__( + self, + img_size: int = 224, + patch_size: int = 8, + embed_dim: int = 256, + depth: int = 2, + num_heads: int = 4, + vocab_size: int = 512, + ) -> None: + """Initializes the BEIT pre-training module. + + Args: + img_size: + Spatial resolution of input images. + patch_size: + Size of each patch. + embed_dim: + Dimension of token embeddings. + depth: + Number of Transformer blocks. + num_heads: + Number of attention heads per block. + vocab_size: + Size of the discrete visual vocabulary. + """ + super().__init__() + self.backbone = BEITEncoder( + img_size=img_size, + patch_size=patch_size, + embed_dim=embed_dim, + depth=depth, + num_heads=num_heads, + ) + self.num_patches = self.backbone.num_patches + self.projection_head = BEiTMIMHead( + embed_dim=embed_dim, + vocab_size=vocab_size, + ) + self.tokenizer = BEiTImageTokenizer(vocab_size=vocab_size) + self.tokenizer.requires_grad_(False) + self.tokenizer.eval() + self.criterion = MaskedImageModelingLoss() + + def training_step( + self, + batch: tuple[torch.Tensor, torch.Tensor], + batch_idx: int, + ) -> torch.Tensor: + """Performs a single training step. + + Args: + batch: + Tuple of (images, targets) where images is a tensor of + shape (B, C, H, W). + batch_idx: + Index of the current batch. + + Returns: + The computed loss value. + """ + images, _ = batch + batch_size = images.shape[0] + + # Generate blockwise mask + bool_masked_pos = self.transform.mask_generator( + batch_size=batch_size, + ).to(device=images.device) + + # Forward pass through encoder + head + enc_out = self.backbone( + x=images, + bool_masked_pos=bool_masked_pos, + ) + patch_features = enc_out["patch_features"] + all_logits = self.projection_head(patch_features=patch_features) + + # Select masked positions + mim_logits = all_logits[bool_masked_pos] + + # Get token targets from frozen tokenizer + with torch.no_grad(): + token_ids = self.tokenizer.tokenize(x=images) + token_targets = token_ids[bool_masked_pos] + + loss = self.criterion(mim_logits, token_targets) + return loss + + def configure_optimizers(self) -> torch.optim.Optimizer: + """Configures the optimizer for training. + + Returns: + AdamW optimizer. + """ + optim = torch.optim.AdamW( + params=self.parameters(), + lr=1.5e-3, + weight_decay=0.05, + ) + return optim + + +model = BEIT() + +transform = BEiTTransform(input_size=224, patch_size=8) +model.transform = transform # Store transform for access in training_step + + +# we ignore object detection annotations by setting target_transform to return 0 +def target_transform(t): + return 0 + + +dataset = torchvision.datasets.FakeData( + size=200, # number of fake samples + image_size=(3, 224, 224), + num_classes=10, + transform=transform, +) + +dataloader = torch.utils.data.DataLoader( + dataset=dataset, + batch_size=16, + shuffle=True, + drop_last=True, + num_workers=2, +) + +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/beit.py b/examples/pytorch_lightning_distributed/beit.py new file mode 100644 index 000000000..7f0c2dbe2 --- /dev/null +++ b/examples/pytorch_lightning_distributed/beit.py @@ -0,0 +1,129 @@ +# 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 lightly.loss import MaskedImageModelingLoss +from lightly.models.modules import BEITEncoder, BEiTImageTokenizer, BEiTMIMHead +from lightly.transforms import BEiTTransform + + +class BEiT(pl.LightningModule): + """BEiT pre-training module for masked image modeling with DDP support. + + Attributes: + backbone: + BEiT Vision Transformer encoder. + projection_head: + MIM head mapping patch features to vocabulary logits. + tokenizer: + Discrete visual tokenizer for generating target token IDs. + criterion: + Cross-entropy loss for masked token prediction. + """ + + def __init__( + self, + img_size: int = 224, + patch_size: int = 8, + embed_dim: int = 256, + depth: int = 2, + num_heads: int = 4, + vocab_size: int = 512, + ) -> None: + super().__init__() + self.backbone = BEITEncoder( + img_size=img_size, + patch_size=patch_size, + embed_dim=embed_dim, + depth=depth, + num_heads=num_heads, + ) + self.num_patches = self.backbone.num_patches + self.projection_head = BEiTMIMHead( + embed_dim=embed_dim, + vocab_size=vocab_size, + ) + self.tokenizer = BEiTImageTokenizer(vocab_size=vocab_size) + self.tokenizer.requires_grad_(False) + self.tokenizer.eval() + self.criterion = MaskedImageModelingLoss() + + def training_step( + self, + batch: tuple[torch.Tensor, torch.Tensor], + batch_idx: int, + ) -> torch.Tensor: + images, _ = batch + batch_size = images.shape[0] + + bool_masked_pos = self.transform.mask_generator( + batch_size=batch_size, + ).to(device=images.device) + + enc_out = self.backbone( + x=images, + bool_masked_pos=bool_masked_pos, + ) + patch_features = enc_out["patch_features"] + all_logits = self.projection_head(patch_features=patch_features) + + mim_logits = all_logits[bool_masked_pos] + + with torch.no_grad(): + token_ids = self.tokenizer.tokenize(x=images) + token_targets = token_ids[bool_masked_pos] + + loss = self.criterion(mim_logits, token_targets) + return loss + + def configure_optimizers(self) -> torch.optim.Optimizer: + optim = torch.optim.AdamW( + params=self.parameters(), + lr=1.5e-3, + weight_decay=0.05, + ) + return optim + + +model = BEiT() + +transform = BEiTTransform(input_size=224, patch_size=8) +model.transform = transform # Store transform for access in training_step + +# Fast +dataset = torchvision.datasets.FakeData( + size=200, + image_size=(3, 224, 224), + num_classes=20, + transform=transform, +) + +dataloader = torch.utils.data.DataLoader( + dataset, + batch_size=16, + shuffle=True, + drop_last=True, + num_workers=2, +) + +accelerator = "gpu" if torch.cuda.is_available() else "cpu" +n_gpus = torch.cuda.device_count() + + +trainer = pl.Trainer( + max_epochs=10, + devices="auto", + accelerator="gpu", + sync_batchnorm=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/loss/__init__.py b/lightly/loss/__init__.py index eeb1c1c48..0350fbe9f 100644 --- a/lightly/loss/__init__.py +++ b/lightly/loss/__init__.py @@ -14,6 +14,7 @@ from lightly.loss.koleo_loss import KoLeoLoss from lightly.loss.lejepa_loss import LeJEPALoss, SIGReg from lightly.loss.macl_loss import MACLLoss +from lightly.loss.mim_loss import MaskedImageModelingLoss from lightly.loss.mmcr_loss import MMCRLoss from lightly.loss.msn_loss import MSNLoss from lightly.loss.negative_cosine_similarity import NegativeCosineSimilarity diff --git a/lightly/loss/mim_loss.py b/lightly/loss/mim_loss.py new file mode 100644 index 000000000..c8cfbe2ca --- /dev/null +++ b/lightly/loss/mim_loss.py @@ -0,0 +1,107 @@ +# Copyright (c) 2024. Lightly AG and its affiliates. +# All Rights Reserved + +from typing import Union + +import torch +import torch.nn.functional as F +from torch import Tensor, nn + + +class MaskedImageModelingLoss(nn.Module): + """Loss for Masked Image Modeling as described in BEIT. + + Computes the cross-entropy between the predicted visual token logits at + masked patch positions and the target discrete visual token indices + produced by the image tokenizer (discrete VAE). + + The objective is: + + L = -E[ sum_{i in M} log p_MIM(z_i | x_M) ] + + where M is the set of masked positions, z_i is the target visual token, + and x_M is the corrupted (masked) image. + + + Attributes: + reduction: + Specifies how to aggregate the per-token losses across the batch. + Must be one of ``"mean"`` (default) or ``"sum"``. + ``"mean"`` divides by the total number of masked tokens across + the batch; ``"sum"`` returns the raw sum. + label_smoothing: + Amount of label smoothing in [0.0, 1.0) applied to the + cross-entropy. Matches the ``label_smoothing`` argument of + :func:`torch.nn.functional.cross_entropy`. Default: 0.0. + + Examples: + >>> import torch + >>> from lightly.loss import MaskedImageModelingLoss + >>> + >>> loss_fn = MaskedImageModelingLoss() + >>> + >>> # mim_logits: predictions at masked positions (N_masked, vocab_size) + >>> # token_targets: ground-truth token ids (N_masked,) + >>> mim_logits = torch.randn(40, 8192) + >>> token_targets = torch.randint(0, 8192, (40,)) + >>> loss = loss_fn(mim_logits, token_targets) + """ + + def __init__( + self, + reduction: str = "mean", + label_smoothing: float = 0.0, + ) -> None: + super().__init__() + + if reduction not in ("mean", "sum"): + raise ValueError( + f"Invalid reduction '{reduction}'. Must be 'mean' or 'sum'." + ) + if not 0.0 <= label_smoothing < 1.0: + raise ValueError( + f"label_smoothing must be in [0, 1), got {label_smoothing}." + ) + + self.reduction = reduction + self.label_smoothing = label_smoothing + + def forward( + self, + mim_logits: Tensor, + token_targets: Tensor, + ) -> Tensor: + """Compute the MIM loss. + + Args: + mim_logits: + Predicted logits for each masked patch position. + Shape: ``(N_masked, vocab_size)`` where ``N_masked`` is the + total number of masked patches across the batch (i.e. + ``batch_size * num_masked_patches_per_image`` when the mask + count is fixed, or the flattened variable-length masked set). + token_targets: + Ground-truth visual token indices from the image tokenizer for + each masked position. Shape: ``(N_masked,)``, dtype ``long``. + + Returns: + Scalar loss tensor. + + Raises: + ValueError: + If ``mim_logits`` and ``token_targets`` have incompatible + leading dimensions. + """ + if mim_logits.shape[0] != token_targets.shape[0]: + raise ValueError( + f"mim_logits has {mim_logits.shape[0]} entries but " + f"token_targets has {token_targets.shape[0]}. " + "Both must have the same number of masked positions." + ) + + return F.cross_entropy( + mim_logits, + token_targets, + reduction=self.reduction, + label_smoothing=self.label_smoothing, + ) diff --git a/lightly/models/modules/__init__.py b/lightly/models/modules/__init__.py index 57341294d..031260dec 100644 --- a/lightly/models/modules/__init__.py +++ b/lightly/models/modules/__init__.py @@ -8,8 +8,10 @@ # Copyright (c) 2021. Lightly AG and its affiliates. # All Rights Reserved +from lightly.models.modules.beit_tokenizer import BEiTImageTokenizer from lightly.models.modules.heads import ( BarlowTwinsProjectionHead, + BEiTMIMHead, BYOLPredictionHead, BYOLProjectionHead, DenseCLProjectionHead, @@ -29,6 +31,7 @@ SwaVProjectionHead, SwaVPrototypes, ) +from lightly.models.modules.masked_image_modeling import BEITEncoder from lightly.models.modules.nn_memory_bank import NNMemoryBankModule from lightly.utils import dependency as _dependency diff --git a/lightly/models/modules/beit_tokenizer.py b/lightly/models/modules/beit_tokenizer.py new file mode 100644 index 000000000..45b617800 --- /dev/null +++ b/lightly/models/modules/beit_tokenizer.py @@ -0,0 +1,513 @@ +# Copyright (c) 2024. Lightly AG and its affiliates. +# All Rights Reserved + +# The image tokenizer follows the discrete VAE architecture from DALL-E [0], +# which is the tokenizer used in the BEIT paper [1]. +# +# - [0] Reference implementation: https://github.com/openai/DALL-E +# - [1]: BEIT, 2021, https://arxiv.org/abs/2106.08254 + +from __future__ import annotations + +from collections import OrderedDict +from functools import partial +from typing import Tuple + +import torch +import torch.nn as nn +import torch.nn.functional as F + + +class _EncoderBlock(nn.Module): + """Residual block used inside the DALL-E encoder. + + Uses a bottleneck residual path with four convolutions and a learned + post-gain of ``1 / n_layers²`` to stabilise deep networks. + """ + + def __init__(self, n_in: int, n_out: int, n_layers: int) -> None: + super().__init__() + n_hid = n_out // 4 + self.post_gain = 1.0 / (n_layers**2) + self.id_path = nn.Conv2d(n_in, n_out, 1) if n_in != n_out else nn.Identity() + self.res_path = nn.Sequential( + OrderedDict( + [ + ("relu_1", nn.ReLU()), + ("conv_1", nn.Conv2d(n_in, n_hid, 3, padding=1)), + ("relu_2", nn.ReLU()), + ("conv_2", nn.Conv2d(n_hid, n_hid, 3, padding=1)), + ("relu_3", nn.ReLU()), + ("conv_3", nn.Conv2d(n_hid, n_hid, 3, padding=1)), + ("relu_4", nn.ReLU()), + ("conv_4", nn.Conv2d(n_hid, n_out, 1)), + ] + ) + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + result: torch.Tensor = self.id_path(x) + self.post_gain * self.res_path(x) + return result + + +class _DecoderBlock(nn.Module): + """Residual block used inside the DALL-E decoder. + + Mirrors ``_EncoderBlock`` but the first bottleneck conv uses a 1×1 kernel + (matching the DALL-E reference implementation). + """ + + def __init__(self, n_in: int, n_out: int, n_layers: int) -> None: + super().__init__() + n_hid = n_out // 4 + self.post_gain = 1.0 / (n_layers**2) + self.id_path = nn.Conv2d(n_in, n_out, 1) if n_in != n_out else nn.Identity() + self.res_path = nn.Sequential( + OrderedDict( + [ + ("relu_1", nn.ReLU()), + ("conv_1", nn.Conv2d(n_in, n_hid, 1)), # 1×1 in decoder + ("relu_2", nn.ReLU()), + ("conv_2", nn.Conv2d(n_hid, n_hid, 3, padding=1)), + ("relu_3", nn.ReLU()), + ("conv_3", nn.Conv2d(n_hid, n_hid, 3, padding=1)), + ("relu_4", nn.ReLU()), + ("conv_4", nn.Conv2d(n_hid, n_out, 3, padding=1)), + ] + ) + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + result: torch.Tensor = self.id_path(x) + self.post_gain * self.res_path(x) + return result + + +class _DALLEEncoder(nn.Module): + """DALL-E discrete VAE encoder (q_φ(z|x)). + + Four groups of residual blocks with max-pool downsampling between groups + (×8 total spatial reduction), followed by a 1×1 conv to ``vocab_size`` + logits. Architecture exactly matches the OpenAI DALL-E reference. + + Args: + n_hid: + Base number of hidden channels (default 256). Each group doubles + this: 1×, 2×, 4×, 8×. + n_blk_per_group: + Number of residual blocks per group (default 2). + input_channels: + Number of input image channels (default 3). + vocab_size: + Codebook / vocabulary size (default 8192). + """ + + group_count: int = 4 + + def __init__( + self, + n_hid: int = 256, + n_blk_per_group: int = 2, + input_channels: int = 3, + vocab_size: int = 8192, + ) -> None: + super().__init__() + self.vocab_size = vocab_size + + blk_range = range(n_blk_per_group) + n_layers = self.group_count * n_blk_per_group + make_blk = partial(_EncoderBlock, n_layers=n_layers) + + self.blocks = nn.Sequential( + OrderedDict( + [ + ("input", nn.Conv2d(input_channels, 1 * n_hid, 7, padding=3)), + ( + "group_1", + nn.Sequential( + OrderedDict( + [ + *[ + ( + f"block_{i + 1}", + make_blk(1 * n_hid, 1 * n_hid), + ) + for i in blk_range + ], + ("pool", nn.MaxPool2d(kernel_size=2)), + ] + ) + ), + ), + ( + "group_2", + nn.Sequential( + OrderedDict( + [ + *[ + ( + f"block_{i + 1}", + make_blk( + 1 * n_hid if i == 0 else 2 * n_hid, + 2 * n_hid, + ), + ) + for i in blk_range + ], + ("pool", nn.MaxPool2d(kernel_size=2)), + ] + ) + ), + ), + ( + "group_3", + nn.Sequential( + OrderedDict( + [ + *[ + ( + f"block_{i + 1}", + make_blk( + 2 * n_hid if i == 0 else 4 * n_hid, + 4 * n_hid, + ), + ) + for i in blk_range + ], + ("pool", nn.MaxPool2d(kernel_size=2)), + ] + ) + ), + ), + ( + "group_4", + nn.Sequential( + OrderedDict( + [ + *[ + ( + f"block_{i + 1}", + make_blk( + 4 * n_hid if i == 0 else 8 * n_hid, + 8 * n_hid, + ), + ) + for i in blk_range + ], + ] + ) + ), + ), + ( + "output", + nn.Sequential( + OrderedDict( + [ + ("relu", nn.ReLU()), + ("conv", nn.Conv2d(8 * n_hid, vocab_size, 1)), + ] + ) + ), + ), + ] + ) + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Encodes an image into logits over the visual vocabulary. + + Args: + x: + Input image tensor of shape ``(B, C, H, W)``, dtype + ``torch.float32``. + + Returns: + Logits tensor of shape ``(B, vocab_size, H/8, W/8)``. + """ + if x.ndim != 4: + raise ValueError(f"input shape {x.shape} is not 4d") + if x.dtype != torch.float32: + raise ValueError("input must have dtype torch.float32") + result: torch.Tensor = self.blocks(x) + return result + + +class _DALLEDecoder(nn.Module): + """DALL-E discrete VAE decoder (p_ψ(x|z)). + + Four groups of residual blocks with nearest-neighbour upsampling between + groups (×8 total spatial increase), followed by a 1×1 conv to + ``2 * output_channels``. The output is split into mean and log-variance + over colour channels. Architecture matches the OpenAI DALL-E reference. + + Args: + n_init: + Initial projection width from the vocabulary embedding (default 128). + n_hid: + Base number of hidden channels (default 256). + n_blk_per_group: + Number of residual blocks per group (default 2). + output_channels: + Number of output image channels (default 3). + vocab_size: + Codebook / vocabulary size (default 8192). + """ + + group_count: int = 4 + + def __init__( + self, + n_init: int = 128, + n_hid: int = 256, + n_blk_per_group: int = 2, + output_channels: int = 3, + vocab_size: int = 8192, + ) -> None: + super().__init__() + self.vocab_size = vocab_size + self.output_channels = output_channels + + blk_range = range(n_blk_per_group) + n_layers = self.group_count * n_blk_per_group + make_blk = partial(_DecoderBlock, n_layers=n_layers) + + self.blocks = nn.Sequential( + OrderedDict( + [ + ("input", nn.Conv2d(vocab_size, n_init, 1)), + ( + "group_1", + nn.Sequential( + OrderedDict( + [ + *[ + ( + f"block_{i + 1}", + make_blk( + n_init if i == 0 else 8 * n_hid, + 8 * n_hid, + ), + ) + for i in blk_range + ], + ( + "upsample", + nn.Upsample(scale_factor=2, mode="nearest"), + ), + ] + ) + ), + ), + ( + "group_2", + nn.Sequential( + OrderedDict( + [ + *[ + ( + f"block_{i + 1}", + make_blk( + 8 * n_hid if i == 0 else 4 * n_hid, + 4 * n_hid, + ), + ) + for i in blk_range + ], + ( + "upsample", + nn.Upsample(scale_factor=2, mode="nearest"), + ), + ] + ) + ), + ), + ( + "group_3", + nn.Sequential( + OrderedDict( + [ + *[ + ( + f"block_{i + 1}", + make_blk( + 4 * n_hid if i == 0 else 2 * n_hid, + 2 * n_hid, + ), + ) + for i in blk_range + ], + ( + "upsample", + nn.Upsample(scale_factor=2, mode="nearest"), + ), + ] + ) + ), + ), + ( + "group_4", + nn.Sequential( + OrderedDict( + [ + *[ + ( + f"block_{i + 1}", + make_blk( + 2 * n_hid if i == 0 else 1 * n_hid, + 1 * n_hid, + ), + ) + for i in blk_range + ], + ] + ) + ), + ), + ( + "output", + nn.Sequential( + OrderedDict( + [ + ("relu", nn.ReLU()), + ( + "conv", + nn.Conv2d(1 * n_hid, 2 * output_channels, 1), + ), + ] + ) + ), + ), + ] + ) + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Decodes a soft token map into a reconstructed image. + + Args: + x: + Soft one-hot token map of shape ``(B, vocab_size, h, w)``, + dtype ``torch.float32``. + + Returns: + Reconstructed image tensor of shape ``(B, C, H, W)`` where + ``C = output_channels``. The decoder output has + ``2 * output_channels`` channels; only the first half (the mean) + is returned — the second half (log-variance) is discarded during + pre-training, consistent with the DALL-E reference. + """ + if x.ndim != 4: + raise ValueError(f"input shape {x.shape} is not 4d") + if x.shape[1] != self.vocab_size: + raise ValueError( + f"input has {x.shape[1]} channels but model built for {self.vocab_size}" + ) + if x.dtype != torch.float32: + raise ValueError("input must have dtype torch.float32") + out: torch.Tensor = self.blocks(x) + return out[:, : self.output_channels] + + +class BEiTImageTokenizer(nn.Module): + """Discrete VAE image tokenizer following the DALL-E architecture . + + Used in BEIT to convert images into discrete visual tokens that serve + as pre-training targets for masked image modelling. The tokenizer is + pre-trained separately (Stage 1) and kept **frozen** during BEIT + pre-training (Stage 2). + + The encoder down-samples by 8× so a 224×224 image becomes a + 14×14 grid of 8192-way token logits. + + + Attributes: + vocab_size: + Size of the visual codebook (default 8192). + temperature: + Gumbel-softmax temperature used during tokenizer training + (default 1.0). Has no effect when calling ``tokenize()``. + + Examples: + >>> import torch + >>> from lightly.models.modules import BEiTImageTokenizer + >>> + >>> tokenizer = BEiTImageTokenizer() + >>> images = torch.randn(2, 3, 224, 224) + >>> + >>> # During BEIT pre-training: get hard token ids (frozen tokenizer) + >>> token_ids = tokenizer.tokenize(images) # (2, 196) + >>> + >>> # During tokenizer training: get logits + reconstruction + >>> logits, recon = tokenizer(images) # (2, 8192, 28, 28), (2, 3, 224, 224) + """ + + def __init__( + self, + input_channels: int = 3, + vocab_size: int = 8192, + n_hid: int = 256, + n_blk_per_group: int = 2, + n_init: int = 128, + temperature: float = 1.0, + ) -> None: + super().__init__() + self.vocab_size = vocab_size + self.temperature = temperature + + self.encoder = _DALLEEncoder( + n_hid=n_hid, + n_blk_per_group=n_blk_per_group, + input_channels=input_channels, + vocab_size=vocab_size, + ) + self.decoder = _DALLEDecoder( + n_init=n_init, + n_hid=n_hid, + n_blk_per_group=n_blk_per_group, + output_channels=input_channels, + vocab_size=vocab_size, + ) + + @torch.no_grad() + def tokenize(self, x: torch.Tensor) -> torch.Tensor: + """Converts images into discrete visual token indices. + + Called with ``torch.no_grad()`` — the tokenizer is frozen during + BEIT pre-training. + + Args: + x: + Input image tensor of shape ``(B, C, H, W)``, + dtype ``torch.float32``. + + Returns: + Token indices of shape ``(B, h*w)`` as a ``torch.long`` tensor, + where ``h = H / 8`` and ``w = W / 8``. + """ + logits: torch.Tensor = self.encoder(x) + token_ids = logits.argmax(dim=1) + B, h, w = token_ids.shape + result: torch.Tensor = token_ids.view(B, h * w) + return result + + def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + """Encodes and decodes an image for tokenizer training. + + Uses Gumbel-softmax to allow gradients to flow through the discrete + bottleneck . + + Args: + x: + Input image tensor of shape ``(B, C, H, W)``, + dtype ``torch.float32``. + + Returns: + A tuple of: + - **logits** ``(B, vocab_size, h, w)`` — encoder output, used + to compute the codebook / reconstruction loss. + - **recon** ``(B, C, H, W)`` — pixel-space reconstruction + from the decoder. + """ + logits: torch.Tensor = self.encoder(x) + soft_one_hot = F.gumbel_softmax(logits, tau=self.temperature, dim=1, hard=False) + recon: torch.Tensor = self.decoder(soft_one_hot) + return logits, recon diff --git a/lightly/models/modules/heads.py b/lightly/models/modules/heads.py index 8c746906d..0f628796b 100644 --- a/lightly/models/modules/heads.py +++ b/lightly/models/modules/heads.py @@ -1031,3 +1031,46 @@ def __init__( (hidden_dim, output_dim, None, None), ] ) + + +class BEiTMIMHead(nn.Module): + """Masked Image Modelling prediction head. + A single linear layer that maps each patch encoding to a + distribution over the visual vocabulary. + + Attributes: + head: + Linear projection from embedding dimension to vocabulary size. + """ + + def __init__( + self, + embed_dim: int = 768, + vocab_size: int = 8192, + ) -> None: + """Initializes BEiTMIMHead. + + Args: + embed_dim: + Dimension of the input patch representations. + vocab_size: + Size of the visual vocabulary (number of discrete tokens). + """ + super().__init__() + self.head = nn.Linear( + in_features=embed_dim, + out_features=vocab_size, + ) + + def forward(self, patch_features: torch.Tensor) -> torch.Tensor: + """Projects patch features to vocabulary logits. + + Args: + patch_features: + Patch representations of shape (B, N, D). + + Returns: + Logits of shape (B, N, vocab_size). + """ + result: torch.Tensor = self.head(patch_features) + return result diff --git a/lightly/models/modules/masked_image_modeling.py b/lightly/models/modules/masked_image_modeling.py new file mode 100644 index 000000000..3ca645c86 --- /dev/null +++ b/lightly/models/modules/masked_image_modeling.py @@ -0,0 +1,885 @@ +# Copyright (c) 2024. Lightly AG and its affiliates. +# All Rights Reserved + + +from __future__ import annotations + +import math +from typing import Optional, Tuple, cast + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from lightly.models.utils import drop_path + + +class DropPath(nn.Module): + """Drop paths (Stochastic Depth) per sample. + + Attributes: + drop_prob: + Probability of dropping a path. If 0, the module is a no-op. + """ + + def __init__(self, drop_prob: float = 0.0) -> None: + """Initializes DropPath. + + Args: + drop_prob: + Probability of dropping a path. Must be in [0, 1). + """ + super().__init__() + self.drop_prob = drop_prob + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Applies stochastic depth to the input tensor. + + Args: + x: + Input tensor of arbitrary shape. + + Returns: + The input tensor, possibly with some samples zeroed out. + """ + result: torch.Tensor = drop_path(x, self.drop_prob, self.training) + return result + + +class PatchEmbed(nn.Module): + """Image to patch embedding. + + Splits an image into non-overlapping patches and projects each patch + into an embedding vector via a convolution. + + Attributes: + img_size: + Spatial resolution of the expected input image as (H, W). + patch_size: + Size of each patch as (patch_h, patch_w). + patch_shape: + Grid shape (num_patches_h, num_patches_w). + num_patches: + Total number of patches, i.e. patch_shape[0] * patch_shape[1]. + proj: + Convolution layer that performs the patch embedding. + """ + + def __init__( + self, + img_size: int | Tuple[int, int] = 224, + patch_size: int | Tuple[int, int] = 16, + in_channels: int = 3, + embed_dim: int = 768, + ) -> None: + """Initializes PatchEmbed. + + Args: + img_size: + Spatial resolution of the input image. Can be a single int + for square images or a tuple (H, W). + patch_size: + Size of each patch. Can be a single int or a tuple. + in_channels: + Number of input channels (e.g. 3 for RGB). + embed_dim: + Dimension of the output embedding vectors. + """ + super().__init__() + img_size_tuple: Tuple[int, int] = ( + (img_size, img_size) if isinstance(img_size, int) else img_size + ) + patch_size_tuple: Tuple[int, int] = ( + (patch_size, patch_size) if isinstance(patch_size, int) else patch_size + ) + self.img_size: Tuple[int, int] = img_size_tuple + self.patch_size: Tuple[int, int] = patch_size_tuple + self.patch_shape: Tuple[int, int] = ( + img_size_tuple[0] // patch_size_tuple[0], + img_size_tuple[1] // patch_size_tuple[1], + ) + self.num_patches: int = self.patch_shape[0] * self.patch_shape[1] + + self.proj = nn.Conv2d( + in_channels=in_channels, + out_channels=embed_dim, + kernel_size=patch_size_tuple, + stride=patch_size_tuple, + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Projects image to patch embeddings. + + Args: + x: + Input image tensor of shape (B, C, H, W). + + Returns: + Patch embeddings of shape (B, N, D) where N is the number of + patches and D is the embedding dimension. + + Raises: + AssertionError: + If the input spatial dimensions do not match img_size. + """ + B, C, H, W = x.shape + assert H == self.img_size[0] and W == self.img_size[1], ( + f"Input image size ({H}*{W}) doesn't match model " + f"({self.img_size[0]}*{self.img_size[1]})." + ) + x = self.proj(x).flatten(start_dim=2).transpose(dim0=1, dim1=2) + return x + + +class Attention(nn.Module): + """Multi-head self-attention with optional relative position bias. + + This implementation follows the original BEIT / BERT design where + query and value have separate bias terms while key has no bias. + + Attributes: + num_heads: + Number of attention heads. + scale: + Scaling factor for query-key dot products (1/sqrt(head_dim)). + qkv: + Linear projection for Q, K, V (without bias). + q_bias: + Learnable bias for the query projection. None if qkv_bias=False. + v_bias: + Learnable bias for the value projection. None if qkv_bias=False. + relative_position_bias_table: + Learnable relative position bias table. None if window_size is None. + relative_position_index: + Registered buffer holding the index mapping for relative positions. + attn_drop: + Dropout layer applied to attention weights. + proj: + Linear projection mapping concatenated heads back to dim. + proj_drop: + Dropout layer applied to the output of proj. + """ + + def __init__( + self, + dim: int, + num_heads: int = 12, + qkv_bias: bool = True, + qk_scale: Optional[float] = None, + attn_drop: float = 0.0, + proj_drop: float = 0.0, + window_size: Optional[Tuple[int, int]] = None, + attn_head_dim: Optional[int] = None, + ) -> None: + """Initializes Attention. + + Args: + dim: + Total dimension of the input (and output). + num_heads: + Number of attention heads. + qkv_bias: + If True, add learnable bias to query and value projections. + qk_scale: + Override for the attention scaling factor. If None, uses + 1/sqrt(head_dim). + attn_drop: + Dropout rate applied to attention weights. + proj_drop: + Dropout rate applied to the output projection. + window_size: + If provided, enables per-head relative position bias with + the given (height, width) patch grid shape. + attn_head_dim: + If provided, overrides the computed head dimension. + """ + super().__init__() + self.num_heads = num_heads + head_dim = dim // num_heads + if attn_head_dim is not None: + head_dim = attn_head_dim + all_head_dim = head_dim * self.num_heads + self.scale = qk_scale or (head_dim**-0.5) + + self.qkv = nn.Linear( + in_features=dim, + out_features=all_head_dim * 3, + bias=False, + ) + self.q_bias: Optional[nn.Parameter] + self.v_bias: Optional[nn.Parameter] + if qkv_bias: + self.q_bias = nn.Parameter(torch.zeros(all_head_dim)) + self.v_bias = nn.Parameter(torch.zeros(all_head_dim)) + else: + self.q_bias = None + self.v_bias = None + + self.window_size: Optional[Tuple[int, int]] + self.relative_position_bias_table: Optional[nn.Parameter] + self.relative_position_index: Optional[torch.Tensor] + if window_size is not None: + self.window_size = window_size + self.num_relative_distance = (2 * window_size[0] - 1) * ( + 2 * window_size[1] - 1 + ) + 3 + self.relative_position_bias_table = nn.Parameter( + torch.zeros(self.num_relative_distance, num_heads) + ) + coords_h = torch.arange(window_size[0]) + coords_w = torch.arange(window_size[1]) + coords = torch.stack(torch.meshgrid([coords_h, coords_w], indexing="ij")) + coords_flatten = torch.flatten(coords, start_dim=1) + relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :] + relative_coords = relative_coords.permute(1, 2, 0).contiguous() + relative_coords[:, :, 0] += window_size[0] - 1 + relative_coords[:, :, 1] += window_size[1] - 1 + relative_coords[:, :, 0] *= 2 * window_size[1] - 1 + relative_position_index = torch.zeros( + size=(window_size[0] * window_size[1] + 1,) * 2, + dtype=relative_coords.dtype, + ) + relative_position_index[1:, 1:] = relative_coords.sum(dim=-1) + relative_position_index[0, 0:] = self.num_relative_distance - 3 + relative_position_index[0:, 0] = self.num_relative_distance - 2 + relative_position_index[0, 0] = self.num_relative_distance - 1 + self.register_buffer( + name="relative_position_index", + tensor=relative_position_index, + ) + else: + self.window_size = None + self.relative_position_bias_table = None + self.relative_position_index = None + + self.attn_drop = nn.Dropout(p=attn_drop) + self.proj = nn.Linear(in_features=all_head_dim, out_features=dim) + self.proj_drop = nn.Dropout(p=proj_drop) + + def forward( + self, + x: torch.Tensor, + rel_pos_bias: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """Computes multi-head self-attention. + + Args: + x: + Input tensor of shape (B, N, C). + rel_pos_bias: + Optional shared relative position bias of shape + (num_heads, N, N) to add to the attention map. + + Returns: + Output tensor of shape (B, N, C). + """ + B, N, C = x.shape + + qkv_bias = None + if self.q_bias is not None: + assert self.v_bias is not None + qkv_bias = torch.cat( + ( + self.q_bias, + torch.zeros_like(self.v_bias, requires_grad=False), + self.v_bias, + ) + ) + + qkv = F.linear(input=x, weight=self.qkv.weight, bias=qkv_bias) + qkv = qkv.reshape(B, N, 3, self.num_heads, -1).permute(2, 0, 3, 1, 4) + q, k, v = qkv[0], qkv[1], qkv[2] + + q = q * self.scale + attn = q @ k.transpose(dim0=-2, dim1=-1) + + if self.relative_position_bias_table is not None: + assert self.relative_position_index is not None + assert self.window_size is not None + relative_position_bias = self.relative_position_bias_table[ + self.relative_position_index.view(-1) + ].view( + self.window_size[0] * self.window_size[1] + 1, + self.window_size[0] * self.window_size[1] + 1, + -1, + ) + relative_position_bias = relative_position_bias.permute( + 2, 0, 1 + ).contiguous() + attn = attn + relative_position_bias.unsqueeze(dim=0) + + if rel_pos_bias is not None: + attn = attn + rel_pos_bias + + attn = attn.softmax(dim=-1) + attn = self.attn_drop(attn) + + x = (attn @ v).transpose(dim0=1, dim1=2).reshape(B, N, -1) + x = self.proj(x) + x = self.proj_drop(x) + return x + + +class MLP(nn.Module): + """Feed-forward network (MLP) inside each Transformer block. + + Follows the original BERT implementation: dropout is applied only + after the second linear layer, not between activation and output. + + Attributes: + fc1: + First linear layer projecting from dim to hidden_dim. + act: + Activation function (GELU). + fc2: + Second linear layer projecting from hidden_dim back to dim. + drop: + Dropout layer applied after fc2. + """ + + def __init__( + self, + in_features: int, + hidden_features: Optional[int] = None, + out_features: Optional[int] = None, + drop: float = 0.0, + ) -> None: + """Initializes MLP. + + Args: + in_features: + Dimension of the input. + hidden_features: + Dimension of the hidden layer. If None, defaults to + in_features. + out_features: + Dimension of the output. If None, defaults to in_features. + drop: + Dropout rate applied after the second linear layer. + """ + super().__init__() + hidden_features = hidden_features or in_features + out_features = out_features or in_features + self.fc1 = nn.Linear( + in_features=in_features, + out_features=hidden_features, + ) + self.act = nn.GELU() + self.fc2 = nn.Linear( + in_features=hidden_features, + out_features=out_features, + ) + self.drop = nn.Dropout(p=drop) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Forward pass through the MLP. + + Args: + x: + Input tensor of shape (B, N, in_features). + + Returns: + Output tensor of shape (B, N, out_features). + """ + x = self.fc1(x) + x = self.act(x) + x = self.fc2(x) + x = self.drop(x) + return x + + +class TransformerBlock(nn.Module): + """Single Transformer encoder block. + + Supports optional LayerScale (gamma) for training stability, + and optional relative position bias. + + Attributes: + norm1: + LayerNorm applied before the attention sub-layer. + attn: + Multi-head self-attention module. + drop_path: + Stochastic depth module. + norm2: + LayerNorm applied before the MLP sub-layer. + mlp: + Feed-forward network. + gamma_1: + LayerScale parameter for the attention branch. None if + init_values is None or 0. + gamma_2: + LayerScale parameter for the MLP branch. None if + init_values is None or 0. + layer_idx: + 1-based index of this block in the encoder stack. + """ + + def __init__( + self, + dim: int, + num_heads: int, + mlp_ratio: float = 4.0, + qkv_bias: bool = True, + qk_scale: Optional[float] = None, + drop: float = 0.0, + attn_drop: float = 0.0, + drop_path_rate: float = 0.0, + layer_idx: int = 0, + init_values: Optional[float] = None, + window_size: Optional[Tuple[int, int]] = None, + attn_head_dim: Optional[int] = None, + ) -> None: + """Initializes TransformerBlock. + + Args: + dim: + Dimension of the input and output. + num_heads: + Number of attention heads. + mlp_ratio: + Ratio of MLP hidden dimension to input dimension. + qkv_bias: + If True, enable bias in query and value projections. + qk_scale: + Override for attention scaling factor. + drop: + Dropout rate for MLP and attention output projections. + attn_drop: + Dropout rate for attention weights. + drop_path_rate: + Stochastic depth rate for this block. + layer_idx: + 1-based index of this block (used for weight rescaling). + init_values: + If provided and > 0, enables LayerScale with this initial + value. + window_size: + If provided, enables per-head relative position bias. + attn_head_dim: + If provided, overrides the computed attention head dimension. + """ + super().__init__() + self.norm1 = nn.LayerNorm(normalized_shape=dim) + self.attn = Attention( + dim=dim, + num_heads=num_heads, + qkv_bias=qkv_bias, + qk_scale=qk_scale, + attn_drop=attn_drop, + proj_drop=drop, + window_size=window_size, + attn_head_dim=attn_head_dim, + ) + self.drop_path = ( + DropPath(drop_prob=drop_path_rate) + if drop_path_rate > 0.0 + else nn.Identity() + ) + self.norm2 = nn.LayerNorm(normalized_shape=dim) + self.mlp = MLP( + in_features=dim, + hidden_features=int(dim * mlp_ratio), + drop=drop, + ) + self.layer_idx = layer_idx + + self.gamma_1: Optional[nn.Parameter] + self.gamma_2: Optional[nn.Parameter] + if init_values is not None and init_values > 0.0: + self.gamma_1 = nn.Parameter( + init_values * torch.ones(dim), + requires_grad=True, + ) + self.gamma_2 = nn.Parameter( + init_values * torch.ones(dim), + requires_grad=True, + ) + else: + self.gamma_1 = None + self.gamma_2 = None + + def forward( + self, + x: torch.Tensor, + rel_pos_bias: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """Forward pass through the block. + + Args: + x: + Input tensor of shape (B, N, dim). + rel_pos_bias: + Optional shared relative position bias. + + Returns: + Output tensor of shape (B, N, dim). + """ + if self.gamma_1 is None: + x = x + self.drop_path(self.attn(self.norm1(x), rel_pos_bias=rel_pos_bias)) + x = x + self.drop_path(self.mlp(self.norm2(x))) + else: + x = x + self.drop_path( + self.gamma_1 * self.attn(self.norm1(x), rel_pos_bias=rel_pos_bias) + ) + x = x + self.drop_path(self.gamma_2 * self.mlp(self.norm2(x))) + return x + + +class RelativePositionBias(nn.Module): + """Shared relative position bias across all blocks. + + Computes a learnable relative position bias table that is shared + by all attention heads in all blocks. + + Attributes: + window_size: + Patch grid shape (height, width). + num_relative_distance: + Number of unique relative distances plus 3 for CLS token + interactions. + relative_position_bias_table: + Learnable table mapping relative position indices to bias + values per head. + relative_position_index: + Registered buffer mapping 2D relative positions to flat + indices. + """ + + def __init__( + self, + window_size: Tuple[int, int], + num_heads: int, + ) -> None: + """Initializes RelativePositionBias. + + Args: + window_size: + Patch grid shape (num_patches_h, num_patches_w). + num_heads: + Number of attention heads. + """ + super().__init__() + self.window_size = window_size + self.num_relative_distance = (2 * window_size[0] - 1) * ( + 2 * window_size[1] - 1 + ) + 3 + self.relative_position_bias_table: nn.Parameter = nn.Parameter( + torch.zeros(self.num_relative_distance, num_heads) + ) + + coords_h = torch.arange(window_size[0]) + coords_w = torch.arange(window_size[1]) + coords = torch.stack(torch.meshgrid([coords_h, coords_w], indexing="ij")) + coords_flatten = torch.flatten(coords, start_dim=1) + relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :] + relative_coords = relative_coords.permute(1, 2, 0).contiguous() + relative_coords[:, :, 0] += window_size[0] - 1 + relative_coords[:, :, 1] += window_size[1] - 1 + relative_coords[:, :, 0] *= 2 * window_size[1] - 1 + relative_position_index = torch.zeros( + size=(window_size[0] * window_size[1] + 1,) * 2, + dtype=relative_coords.dtype, + ) + relative_position_index[1:, 1:] = relative_coords.sum(dim=-1) + relative_position_index[0, 0:] = self.num_relative_distance - 3 + relative_position_index[0:, 0] = self.num_relative_distance - 2 + relative_position_index[0, 0] = self.num_relative_distance - 1 + self.relative_position_index: torch.Tensor + self.register_buffer( + name="relative_position_index", + tensor=relative_position_index, + ) + + def forward(self) -> torch.Tensor: + """Computes the shared relative position bias. + + Returns: + Relative position bias of shape + (num_heads, window_size[0]*window_size[1]+1, + window_size[0]*window_size[1]+1). + """ + relative_position_bias = self.relative_position_bias_table[ + self.relative_position_index.view(-1) + ].view( + self.window_size[0] * self.window_size[1] + 1, + self.window_size[0] * self.window_size[1] + 1, + -1, + ) + return relative_position_bias.permute(2, 0, 1).contiguous() + + +class BEITEncoder(nn.Module): + """BEIT Vision Transformer encoder. + + The [S] (CLS) token is prepended; 1-D learnable position embeddings + are added. A learnable mask embedding [M] is used to replace masked + patches. Supports optional relative position bias and LayerScale. + + Attributes: + embed_dim: + Dimension of the token embeddings. + patch_size: + Size of each image patch. + num_patches: + Total number of patches per image. + patch_shape: + Grid shape of patches (height, width). + patch_embed: + Patch embedding module. + cls_token: + Learnable [S] (CLS) token prepended to the sequence. + mask_token: + Learnable [M] token used to replace masked patches. + pos_embed: + Learnable 1-D position embeddings. None if use_abs_pos_emb + is False. + pos_drop: + Dropout applied after adding position embeddings. + rel_pos_bias: + Shared relative position bias module. None if disabled. + blocks: + List of Transformer blocks. + norm: + Final LayerNorm. + """ + + def __init__( + self, + img_size: int = 224, + patch_size: int = 16, + in_channels: int = 3, + embed_dim: int = 768, + depth: int = 12, + num_heads: int = 12, + mlp_ratio: float = 4.0, + qkv_bias: bool = True, + qk_scale: Optional[float] = None, + drop_rate: float = 0.0, + attn_drop_rate: float = 0.0, + drop_path_rate: float = 0.1, + init_values: Optional[float] = None, + use_abs_pos_emb: bool = True, + use_rel_pos_bias: bool = False, + use_shared_rel_pos_bias: bool = False, + attn_head_dim: Optional[int] = None, + init_std: float = 0.02, + ) -> None: + """Initializes BEITEncoder. + + Args: + img_size: + Spatial resolution of input images. + patch_size: + Size of each patch. + in_channels: + Number of input channels. + embed_dim: + Dimension of token embeddings. + depth: + Number of Transformer blocks. + num_heads: + Number of attention heads per block. + mlp_ratio: + Ratio of MLP hidden dimension to embedding dimension. + qkv_bias: + If True, enable bias in query and value projections. + qk_scale: + Override for attention scaling factor. + drop_rate: + Dropout rate for MLP and attention output. + attn_drop_rate: + Dropout rate for attention weights. + drop_path_rate: + Maximum stochastic depth rate (linearly decayed per block). + init_values: + If provided and > 0, enables LayerScale with this initial + value. + use_abs_pos_emb: + If True, use learnable absolute position embeddings. + use_rel_pos_bias: + If True, use per-block relative position bias. + use_shared_rel_pos_bias: + If True, use a single shared relative position bias across + all blocks. + attn_head_dim: + If provided, overrides the computed attention head dimension. + init_std: + Standard deviation for truncated normal initialization. + """ + super().__init__() + self.embed_dim = embed_dim + self.patch_size = patch_size + self.num_patches = (img_size // patch_size) ** 2 + self.init_std = init_std + + self.patch_embed = PatchEmbed( + img_size=img_size, + patch_size=patch_size, + in_channels=in_channels, + embed_dim=embed_dim, + ) + self.patch_shape = self.patch_embed.patch_shape + + self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim)) + self.mask_token = nn.Parameter(torch.zeros(1, 1, embed_dim)) + self.pos_embed: Optional[nn.Parameter] + if use_abs_pos_emb: + self.pos_embed = nn.Parameter( + torch.zeros(1, self.num_patches + 1, embed_dim) + ) + else: + self.pos_embed = None + self.pos_drop = nn.Dropout(p=drop_rate) + + self.rel_pos_bias: Optional[RelativePositionBias] + if use_shared_rel_pos_bias: + self.rel_pos_bias = RelativePositionBias( + window_size=self.patch_shape, + num_heads=num_heads, + ) + else: + self.rel_pos_bias = None + + dpr = [ + x.item() + for x in torch.linspace( + start=0, + end=drop_path_rate, + steps=depth, + ) + ] + + self.blocks = nn.ModuleList( + [ + TransformerBlock( + dim=embed_dim, + num_heads=num_heads, + mlp_ratio=mlp_ratio, + qkv_bias=qkv_bias, + qk_scale=qk_scale, + drop=drop_rate, + attn_drop=attn_drop_rate, + drop_path_rate=dpr[i], + layer_idx=i + 1, + init_values=init_values, + window_size=self.patch_shape if use_rel_pos_bias else None, + attn_head_dim=attn_head_dim, + ) + for i in range(depth) + ] + ) + self.norm = nn.LayerNorm(normalized_shape=embed_dim) + + self._init_weights() + + def _init_weights(self) -> None: + """Initializes all weights using truncated normal.""" + if self.pos_embed is not None: + nn.init.trunc_normal_(self.pos_embed, std=self.init_std) + nn.init.trunc_normal_(self.cls_token, std=self.init_std) + nn.init.trunc_normal_(self.mask_token, std=self.init_std) + self.apply(self._init_weights_recursive) + self.fix_init_weight() + + def _init_weights_recursive(self, m: nn.Module) -> None: + """Recursive weight initialization applied via nn.Module.apply. + + Args: + m: + Module to initialize. + """ + if isinstance(m, nn.Linear): + nn.init.trunc_normal_(m.weight, std=self.init_std) + if isinstance(m, nn.Linear) and m.bias is not None: + nn.init.constant_(m.bias, val=0) + elif isinstance(m, nn.LayerNorm): + nn.init.constant_(m.bias, val=0) + nn.init.constant_(m.weight, val=1.0) + elif isinstance(m, nn.Conv2d): + nn.init.trunc_normal_(m.weight, std=self.init_std) + if m.bias is not None: + nn.init.constant_(m.bias, val=0) + + def fix_init_weight(self) -> None: + """Rescales output projection weights by layer depth. + + This rescaling is critical for training deep transformers from + scratch, as described in the BEIT paper. + """ + + def rescale(param: torch.Tensor, layer_id: int) -> None: + param.div_(math.sqrt(2.0 * layer_id)) + + for layer_id, layer in enumerate(self.blocks, start=1): + block = cast(TransformerBlock, layer) + rescale(param=block.attn.proj.weight.data, layer_id=layer_id) + rescale(param=block.mlp.fc2.weight.data, layer_id=layer_id) + + def apply_mask( + self, + x: torch.Tensor, + bool_masked_pos: torch.Tensor, + ) -> torch.Tensor: + """Replaces masked patch positions with the learnable mask token. + + Args: + x: + Patch embeddings of shape (B, N, D) without CLS token. + bool_masked_pos: + Boolean mask of shape (B, N) where True indicates a masked + position. + + Returns: + Masked patch embeddings of shape (B, N, D). + """ + mask_tokens = self.mask_token.expand(x.shape[0], x.shape[1], -1) + w = bool_masked_pos.unsqueeze(dim=-1).type_as(mask_tokens) + x = x * (1 - w) + mask_tokens * w + return x + + def forward( + self, + x: torch.Tensor, + bool_masked_pos: Optional[torch.Tensor] = None, + ) -> dict[str, torch.Tensor]: + """Forward pass through the encoder. + + Args: + x: + Input image tensor of shape (B, C, H, W). + bool_masked_pos: + Optional boolean mask of shape (B, N) indicating which + patch positions are masked. If None, no masking is applied. + + Returns: + Dictionary with the following keys: + - 'last_hidden_state': (B, N+1, D) full sequence output + - 'patch_features': (B, N, D) patch tokens only + - 'cls_feature': (B, D) CLS token only + """ + B = x.shape[0] + + tokens = self.patch_embed(x=x) + + if bool_masked_pos is not None: + tokens = self.apply_mask(x=tokens, bool_masked_pos=bool_masked_pos) + + cls_tokens = self.cls_token.expand(B, -1, -1) + tokens = torch.cat([cls_tokens, tokens], dim=1) + + if self.pos_embed is not None: + tokens = tokens + self.pos_embed + tokens = self.pos_drop(tokens) + + rel_pos_bias = self.rel_pos_bias() if self.rel_pos_bias is not None else None + for block in self.blocks: + tokens = block(x=tokens, rel_pos_bias=rel_pos_bias) + + tokens = self.norm(tokens) + + return { + "last_hidden_state": tokens, + "patch_features": tokens[:, 1:], + "cls_feature": tokens[:, 0], + } diff --git a/lightly/models/utils.py b/lightly/models/utils.py index eb6c17533..d3b2bb3ef 100644 --- a/lightly/models/utils.py +++ b/lightly/models/utils.py @@ -1316,3 +1316,31 @@ def apply_masks(x: Tensor, masks: Tensor | list[Tensor]) -> Tensor: mask_keep = m.unsqueeze(-1).repeat(1, 1, x.size(-1)) all_x += [torch.gather(x, dim=1, index=mask_keep)] return torch.cat(all_x, dim=0) + + +def drop_path( + x: torch.Tensor, + drop_prob: float = 0.0, + training: bool = False, +) -> torch.Tensor: + """Applies stochastic depth (per sample). + + Args: + x: + Input tensor. + drop_prob: + Probability of dropping a sample. Defaults to 0.0. + training: + Whether in training mode. Defaults to False. + + Returns: + The input tensor with stochastic depth applied during training, + or the unchanged tensor during evaluation or when drop_prob is 0. + """ + if drop_prob == 0.0 or not training: + return x + keep_prob = 1.0 - drop_prob + shape = (x.shape[0],) + (1,) * (x.ndim - 1) + random_tensor = torch.rand(shape, dtype=x.dtype, device=x.device) + random_tensor = torch.floor(random_tensor + keep_prob) + return x.div(keep_prob) * random_tensor diff --git a/lightly/transforms/__init__.py b/lightly/transforms/__init__.py index 6ded4bc68..c146a6462 100644 --- a/lightly/transforms/__init__.py +++ b/lightly/transforms/__init__.py @@ -10,6 +10,7 @@ from lightly.transforms.aim_transform import AIMTransform from lightly.transforms.amplitude_rescale_transform import AmplitudeRescaleTransform +from lightly.transforms.beit_transform import BEiTTransform, BlockwiseMaskingGenerator from lightly.transforms.byol_transform import ( BYOLTransform, BYOLView1Transform, diff --git a/lightly/transforms/beit_transform.py b/lightly/transforms/beit_transform.py new file mode 100644 index 000000000..850fe455a --- /dev/null +++ b/lightly/transforms/beit_transform.py @@ -0,0 +1,274 @@ +# Copyright (c) 2024. Lightly AG and its affiliates. +# All Rights Reserved + +# Implements the data augmentation and blockwise masking strategy described in +# BEiT: BERT Pre-Training of Image Transformers, 2021 +# https://arxiv.org/abs/2106.08254 + +from __future__ import annotations + +import math +from typing import Tuple + +import torch +import torchvision.transforms as T +from PIL import Image + + +class BlockwiseMaskingGenerator: + """Generates random block-shaped masks. + + At each call a random block (with random aspect ratio) is selected and + added to the mask until the target masking ratio is reached. + + Args: + num_patches_per_side: + Grid side length. Default is 14 (for 224×224 images with 16×16 + patches). + mask_ratio: + Target fraction of patches to mask (default 0.4). + min_block_patches: + Minimum number of patches per block (default 16). + max_block_patches: + Maximum number of patches per block. If ``None``, defaults to + the total number of patches to mask (default ``None``). + aspect_ratio_range: + ``(min_r, max_r)`` range for the block aspect ratio sampled + log-uniformly. Default is ``(0.3, 1/0.3)``. + + Examples: + >>> generator = BlockwiseMaskingGenerator(num_patches_per_side=14) + >>> mask = generator(batch_size=4) # (4, 196) + """ + + def __init__( + self, + num_patches_per_side: int = 14, + mask_ratio: float = 0.4, + min_block_patches: int = 16, + max_block_patches: int | None = None, + aspect_ratio_range: Tuple[float, float] = (0.3, 1.0 / 0.3), + ) -> None: + self.h = num_patches_per_side + self.w = num_patches_per_side + self.N = self.h * self.w + self.target_count = int(math.ceil(mask_ratio * self.N)) + self.min_block = min_block_patches + self.max_block = ( + self.target_count if max_block_patches is None else max_block_patches + ) + self.ar_min, self.ar_max = aspect_ratio_range + + def __repr__(self) -> str: + repr_str = ( + "BlockwiseMaskingGenerator(%d, %d -> [%d ~ %d], max = %d, %.3f ~ %.3f)" + % ( + self.h, + self.w, + self.min_block, + self.max_block, + self.target_count, + math.log(self.ar_min), + math.log(self.ar_max), + ) + ) + return repr_str + + def get_shape(self) -> Tuple[int, int]: + """Return the spatial shape of the mask grid.""" + return self.h, self.w + + def _mask(self, mask: torch.Tensor, max_mask_patches: int) -> int: + """Attempt to add one masked block to the current mask. + + Tries up to 10 times to find a block that adds between 1 and + ``max_mask_patches`` new masked patches. + + Args: + mask: + Current boolean mask of shape ``(h, w)``. + max_mask_patches: + Upper bound on the number of *newly* masked patches this + block may add. + + Returns: + Number of newly masked patches added (0 if no valid block found). + """ + delta = 0 + for _ in range(10): + target_area = ( + torch.empty(1) + .uniform_(self.min_block, max(self.min_block + 1, max_mask_patches + 1)) + .item() + ) + ar = math.exp( + torch.empty(1) + .uniform_(math.log(self.ar_min), math.log(self.ar_max)) + .item() + ) + a = max(1, int(round(math.sqrt(target_area * ar)))) + b = max(1, int(round(math.sqrt(target_area / ar)))) + if a >= self.h or b >= self.w: + continue + t = int(torch.randint(0, self.h - a + 1, (1,)).item()) + l = int(torch.randint(0, self.w - b + 1, (1,)).item()) + + num_masked = int(mask[t : t + a, l : l + b].sum().item()) + new_patches = a * b - num_masked + if 0 < new_patches <= max_mask_patches: + mask[t : t + a, l : l + b] = True + delta = new_patches + break + return delta + + def __call__(self, batch_size: int = 1) -> torch.Tensor: + """Generate a batch of block-shaped boolean masks. + + Args: + batch_size: + Number of masks to generate. + + Returns: + Boolean tensor of shape ``(batch_size, N)`` where ``N = + num_patches_per_side²``. ``True`` indicates a masked position. + """ + masks = [] + for _ in range(batch_size): + mask = torch.zeros(self.h, self.w, dtype=torch.bool) + mask_count = 0 + while mask_count < self.target_count: + max_mask_patches = self.target_count - mask_count + max_mask_patches = min(max_mask_patches, self.max_block) + + delta = self._mask(mask, max_mask_patches) + if delta == 0: + break + mask_count += delta + masks.append(mask.flatten()) + return torch.stack(masks, dim=0) + + +class BEiTTransform: + """Image transform for BEiT pre-training. + + Applies the standard BEiT augmentation pipeline (random resized crop, + horizontal flip, colour jitter, normalisation) and exposes a + :attr:`mask_generator` for producing blockwise boolean masks. + + In the training loop, call ``transform`` on each PIL image to get the + augmented tensor, then call ``transform.mask_generator(batch_size=B)`` + once per batch to obtain the mask that is passed to the encoder. + + Attributes: + mask_generator: + A :class:`BlockwiseMaskingGenerator` instance configured with the + parameters supplied to this transform. + + Args: + input_size: + Spatial size of the input image after cropping (default 224). + patch_size: + Patch size used by the ViT backbone; determines the grid side + length as ``input_size // patch_size`` (default 16). + mask_ratio: + Fraction of patches to mask (default 0.4). + min_block_patches: + Minimum number of patches per masked block (default 16). + max_block_patches: + Maximum number of patches per masked block. If ``None``, defaults + to the total number of patches to mask (default ``None``). + color_jitter: + Colour jitter strength applied to brightness, contrast, saturation + and hue (default 0.4, matching §2.5 of the paper). + scale: + Range of size of the origin size cropped (default ``(0.2, 1.0)``). + The original BEiT paper uses ``(0.08, 1.0)``. + ratio: + Range of aspect ratio of the origin aspect ratio cropped + (default ``(3.0 / 4.0, 4.0 / 3.0)``). + interpolation: + Interpolation mode for resizing. Default is ``'bicubic'``. + Options: ``'bilinear'``, ``'bicubic'``, ``'lanczos'``. + mean: + Channel-wise normalisation mean (default ImageNet mean). + std: + Channel-wise normalisation std (default ImageNet std). + + Examples: + >>> from lightly.transforms import BEiTTransform + >>> transform = BEiTTransform() + >>> + >>> # applied per image in the dataset + >>> image = transform(pil_image) # (3, 224, 224) + >>> + >>> # called once per batch inside the dataloader collation / training loop + >>> mask = transform.mask_generator(batch_size=64) # (64, 196) + """ + + def __init__( + self, + input_size: int = 224, + patch_size: int = 16, + mask_ratio: float = 0.4, + min_block_patches: int = 16, + max_block_patches: int | None = None, + color_jitter: float = 0.4, + scale: Tuple[float, float] = (0.2, 1.0), + ratio: Tuple[float, float] = (3.0 / 4.0, 4.0 / 3.0), + aspect_ratio_range: Tuple[float, float] = (0.3, 1.0 / 0.3), + interpolation: str = "bicubic", + mean: Tuple[float, float, float] = (0.485, 0.456, 0.406), + std: Tuple[float, float, float] = (0.229, 0.224, 0.225), + ) -> None: + num_patches_per_side = input_size // patch_size + + self.mask_generator = BlockwiseMaskingGenerator( + num_patches_per_side=num_patches_per_side, + mask_ratio=mask_ratio, + aspect_ratio_range=aspect_ratio_range, + min_block_patches=min_block_patches, + max_block_patches=max_block_patches, + ) + + # Map string interpolation to PIL constants + interpolation_modes = { + "bilinear": T.InterpolationMode.BILINEAR, + "bicubic": T.InterpolationMode.BICUBIC, + "lanczos": T.InterpolationMode.LANCZOS, + } + interp_mode = interpolation_modes.get( + interpolation, T.InterpolationMode.BICUBIC + ) + + self.transform = T.Compose( + [ + T.RandomResizedCrop( + input_size, + scale=scale, + ratio=ratio, + interpolation=interp_mode, + ), + T.RandomHorizontalFlip(), + T.ColorJitter( + brightness=color_jitter, + contrast=color_jitter, + saturation=color_jitter, + hue=color_jitter * 0.25, + ), + T.ToTensor(), + T.Normalize(mean=mean, std=std), + ] + ) + + def __call__(self, image: Image.Image) -> torch.Tensor: + """Apply the augmentation pipeline to a single PIL image. + + Args: + image: + Input PIL image. + + Returns: + Augmented image tensor of shape ``(C, H, W)``. + """ + result: torch.Tensor = self.transform(image) + return result diff --git a/tests/loss/test_mim_loss.py b/tests/loss/test_mim_loss.py new file mode 100644 index 000000000..5f478ec48 --- /dev/null +++ b/tests/loss/test_mim_loss.py @@ -0,0 +1,91 @@ +import unittest + +import torch + +from lightly.loss import MaskedImageModelingLoss + + +class TestMaskedImageModelingLoss(unittest.TestCase): + def test_forward_pass(self) -> None: + loss_fn = MaskedImageModelingLoss() + for n_masked in [1, 40, 160]: + logits = torch.randn(n_masked, 8192) + targets = torch.randint(0, 8192, (n_masked,)) + loss = loss_fn(logits, targets) + self.assertEqual(loss.shape, torch.Size([])) + self.assertTrue(torch.isfinite(loss)) + + def test_reduction_mean(self) -> None: + loss_fn = MaskedImageModelingLoss(reduction="mean") + logits = torch.randn(80, 8192) + targets = torch.randint(0, 8192, (80,)) + loss = loss_fn(logits, targets) + self.assertEqual(loss.shape, torch.Size([])) + + def test_reduction_sum(self) -> None: + loss_mean = MaskedImageModelingLoss(reduction="mean") + loss_sum = MaskedImageModelingLoss(reduction="sum") + logits = torch.randn(80, 8192) + targets = torch.randint(0, 8192, (80,)) + self.assertAlmostEqual( + loss_sum(logits, targets).item(), + loss_mean(logits, targets).item() * 80, + places=4, + ) + + def test_invalid_reduction_raises(self) -> None: + with self.assertRaises(ValueError): + MaskedImageModelingLoss(reduction="none") + + def test_invalid_label_smoothing_raises(self) -> None: + with self.assertRaises(ValueError): + MaskedImageModelingLoss(label_smoothing=-0.1) + with self.assertRaises(ValueError): + MaskedImageModelingLoss(label_smoothing=1.0) + + def test_shape_mismatch_raises(self) -> None: + loss_fn = MaskedImageModelingLoss() + with self.assertRaises(ValueError): + loss_fn(torch.randn(10, 8192), torch.randint(0, 8192, (9,))) + + def test_label_smoothing(self) -> None: + loss_plain = MaskedImageModelingLoss(label_smoothing=0.0) + loss_smooth = MaskedImageModelingLoss(label_smoothing=0.1) + torch.manual_seed(0) + logits = torch.randn(80, 8192) + targets = torch.randint(0, 8192, (80,)) + self.assertNotAlmostEqual( + loss_plain(logits, targets).item(), + loss_smooth(logits, targets).item(), + places=3, + ) + + def test_gradient_flows(self) -> None: + loss_fn = MaskedImageModelingLoss() + logits = torch.randn(40, 8192, requires_grad=True) + targets = torch.randint(0, 8192, (40,)) + loss = loss_fn(logits, targets) + loss.backward() + self.assertIsNotNone(logits.grad) + assert ( + logits.grad is not None + ) # mypy narrowing; assertIsNotNone above doesn't do this + self.assertEqual(logits.grad.shape, logits.shape) + + def test_perfect_prediction_gives_low_loss(self) -> None: + loss_fn = MaskedImageModelingLoss() + n, v = 40, 8192 + targets = torch.randint(0, v, (n,)) + logits = torch.full((n, v), -100.0) + logits[torch.arange(n), targets] = 100.0 + loss = loss_fn(logits, targets) + self.assertLess(loss.item(), 0.01) + + @unittest.skipUnless(torch.cuda.is_available(), "Cuda not available") + def test_forward_pass_cuda(self) -> None: + loss_fn = MaskedImageModelingLoss() + logits = torch.randn(80, 8192).cuda() + targets = torch.randint(0, 8192, (80,)).cuda() + loss = loss_fn(logits, targets) + self.assertEqual(loss.shape, torch.Size([])) + self.assertTrue(torch.isfinite(loss)) diff --git a/tests/models/test_ModelBEiT.py b/tests/models/test_ModelBEiT.py new file mode 100644 index 000000000..1543f4546 --- /dev/null +++ b/tests/models/test_ModelBEiT.py @@ -0,0 +1,281 @@ +from __future__ import annotations + +import math +import unittest +from typing import Any + +import torch + +from lightly.models.modules import BEITEncoder +from lightly.models.modules.heads import BEiTMIMHead + + +class TestBEITEncoder(unittest.TestCase): + """Tests for the BEITEncoder module.""" + + _EMBED_DIM = 64 + _DEPTH = 2 + _NUM_HEADS = 4 + _IMG_SIZE = 32 + _PATCH_SIZE = 8 + + def _make_encoder(self, **kwargs: Any) -> BEITEncoder: + """Creates a BEITEncoder with default test parameters. + + Args: + **kwargs: + Additional keyword arguments passed to BEITEncoder. + + Returns: + A BEITEncoder instance configured for testing. + """ + return BEITEncoder( + img_size=self._IMG_SIZE, + patch_size=self._PATCH_SIZE, + embed_dim=self._EMBED_DIM, + depth=self._DEPTH, + num_heads=self._NUM_HEADS, + **kwargs, + ) + + def _make_mask(self, batch_size: int, mask_ratio: float = 0.4) -> torch.Tensor: + """Creates a deterministic boolean mask for testing. + + Args: + batch_size: + Number of samples in the batch. + mask_ratio: + Fraction of patches to mask. + + Returns: + Boolean tensor of shape (batch_size, n_patches) with the + first n_masked positions set to True. + """ + n_patches = (self._IMG_SIZE // self._PATCH_SIZE) ** 2 + n_masked = int(n_patches * mask_ratio) + mask = torch.zeros(batch_size, n_patches, dtype=torch.bool) + mask[:, :n_masked] = True + return mask + + def test_output_shapes_no_mask(self) -> None: + """Tests output shapes when no mask is provided.""" + encoder = self._make_encoder() + x = torch.randn(2, 3, self._IMG_SIZE, self._IMG_SIZE) + out = encoder(x=x) + n_patches = (self._IMG_SIZE // self._PATCH_SIZE) ** 2 + + self.assertEqual( + out["last_hidden_state"].shape, + (2, n_patches + 1, self._EMBED_DIM), + ) + self.assertEqual( + out["patch_features"].shape, + (2, n_patches, self._EMBED_DIM), + ) + self.assertEqual( + out["cls_feature"].shape, + (2, self._EMBED_DIM), + ) + + def test_output_shapes_with_mask(self) -> None: + """Tests output shapes when masking is applied.""" + encoder = self._make_encoder() + x = torch.randn(2, 3, self._IMG_SIZE, self._IMG_SIZE) + mask = self._make_mask(batch_size=2) + out = encoder(x=x, bool_masked_pos=mask) + n_patches = (self._IMG_SIZE // self._PATCH_SIZE) ** 2 + + self.assertEqual( + out["patch_features"].shape, + (2, n_patches, self._EMBED_DIM), + ) + + def test_masked_vs_unmasked_outputs_differ(self) -> None: + """Tests that masking changes the output.""" + torch.manual_seed(0) + encoder = self._make_encoder() + encoder.eval() + x = torch.randn(1, 3, self._IMG_SIZE, self._IMG_SIZE) + mask = self._make_mask(batch_size=1) + + with torch.no_grad(): + out_clean = encoder(x=x)["patch_features"] + out_masked = encoder(x=x, bool_masked_pos=mask)["patch_features"] + + self.assertFalse(torch.allclose(out_clean, out_masked)) + + def test_unmasked_positions_change_with_context(self) -> None: + """Tests that unmasked positions are affected by masked positions.""" + torch.manual_seed(0) + encoder = self._make_encoder() + encoder.eval() + x = torch.randn(1, 3, self._IMG_SIZE, self._IMG_SIZE) + mask = self._make_mask(batch_size=1, mask_ratio=0.5) + + with torch.no_grad(): + feats_clean = encoder(x=x)["patch_features"] + feats_masked = encoder(x=x, bool_masked_pos=mask)["patch_features"] + + unmasked = ~mask[0] + self.assertFalse( + torch.allclose(feats_clean[:, unmasked], feats_masked[:, unmasked]) + ) + + def test_gradient_flows_through_encoder(self) -> None: + """Tests that gradients propagate through the encoder.""" + encoder = self._make_encoder() + x = torch.randn(1, 3, self._IMG_SIZE, self._IMG_SIZE, requires_grad=True) + mask = self._make_mask(batch_size=1) + out = encoder(x=x, bool_masked_pos=mask) + out["patch_features"].sum().backward() + + self.assertIsNotNone(x.grad) + + def test_mask_token_is_parameter(self) -> None: + """Tests that mask_token is a learnable parameter.""" + encoder = self._make_encoder() + + self.assertIsInstance(encoder.mask_token, torch.nn.Parameter) + self.assertEqual(encoder.mask_token.shape, (1, 1, self._EMBED_DIM)) + + def test_cls_token_is_parameter(self) -> None: + """Tests that cls_token is a learnable parameter.""" + encoder = self._make_encoder() + + self.assertIsInstance(encoder.cls_token, torch.nn.Parameter) + self.assertEqual(encoder.cls_token.shape, (1, 1, self._EMBED_DIM)) + + def test_consistent_across_batch_sizes(self) -> None: + """Tests that outputs are consistent regardless of batch size.""" + encoder = self._make_encoder() + encoder.eval() + x = torch.randn(1, 3, self._IMG_SIZE, self._IMG_SIZE) + mask_1 = self._make_mask(batch_size=1) + mask_2 = mask_1.expand(2, -1) + + with torch.no_grad(): + out_1 = encoder(x=x, bool_masked_pos=mask_1)["patch_features"] + out_2 = encoder( + x=x.expand(2, -1, -1, -1), + bool_masked_pos=mask_2, + )["patch_features"] + + self.assertTrue(torch.allclose(out_1, out_2[:1], atol=1e-5)) + + def test_relative_position_bias(self) -> None: + """Tests encoder with shared relative position bias enabled.""" + encoder = self._make_encoder(use_shared_rel_pos_bias=True) + x = torch.randn(1, 3, self._IMG_SIZE, self._IMG_SIZE) + out = encoder(x=x) + + self.assertEqual( + out["patch_features"].shape, + (1, (self._IMG_SIZE // self._PATCH_SIZE) ** 2, self._EMBED_DIM), + ) + + def test_no_absolute_pos_emb(self) -> None: + """Tests encoder without absolute position embeddings.""" + encoder = self._make_encoder(use_abs_pos_emb=False) + x = torch.randn(1, 3, self._IMG_SIZE, self._IMG_SIZE) + out = encoder(x=x) + + self.assertIsNone(encoder.pos_embed) + self.assertEqual( + out["patch_features"].shape, + (1, (self._IMG_SIZE // self._PATCH_SIZE) ** 2, self._EMBED_DIM), + ) + + def test_layer_scale_init(self) -> None: + """Tests that LayerScale parameters are created when init_values > 0.""" + encoder = self._make_encoder(init_values=0.1) + + for block in encoder.blocks: + self.assertIsNotNone(block.gamma_1) + self.assertIsNotNone(block.gamma_2) + self.assertEqual(block.gamma_1.shape, (self._EMBED_DIM,)) + self.assertEqual(block.gamma_2.shape, (self._EMBED_DIM,)) + + def test_no_layer_scale_by_default(self) -> None: + """Tests that LayerScale is disabled by default.""" + encoder = self._make_encoder() + + for block in encoder.blocks: + self.assertIsNone(block.gamma_1) + self.assertIsNone(block.gamma_2) + + def test_fix_init_weight_applied(self) -> None: + """Tests that fix_init_weight rescales output projections.""" + encoder = self._make_encoder() + encoder.eval() + + for layer_id, block in enumerate(encoder.blocks, start=1): + expected_scale = 1.0 / math.sqrt(2.0 * layer_id) + # We can't easily test the exact value since init is random, + # but we can verify the method ran without error by checking + # the model produces valid outputs. + x = torch.randn(1, 3, self._IMG_SIZE, self._IMG_SIZE) + out = encoder(x=x) + self.assertTrue(torch.isfinite(out["patch_features"]).all()) + + @unittest.skipUnless(torch.cuda.is_available(), "CUDA not available") + def test_forward_pass_cuda(self) -> None: + """Tests forward pass on CUDA if available.""" + encoder = self._make_encoder().cuda() + x = torch.randn(2, 3, self._IMG_SIZE, self._IMG_SIZE).cuda() + mask = self._make_mask(batch_size=2).cuda() + out = encoder(x=x, bool_masked_pos=mask) + + self.assertTrue(out["patch_features"].is_cuda) + + +class TestMIMHead(unittest.TestCase): + """Tests for the BEiTMIMHead module.""" + + _EMBED_DIM = 64 + _VOCAB_SIZE = 512 + + def _make_head(self) -> BEiTMIMHead: + """Creates a BEiTMIMHead with default test parameters. + + Returns: + A BEiTMIMHead instance configured for testing. + """ + return BEiTMIMHead( + embed_dim=self._EMBED_DIM, + vocab_size=self._VOCAB_SIZE, + ) + + def test_output_shape(self) -> None: + """Tests that output shape matches (B, N, vocab_size).""" + head = self._make_head() + features = torch.randn(2, 16, self._EMBED_DIM) + logits = head(patch_features=features) + + self.assertEqual(logits.shape, (2, 16, self._VOCAB_SIZE)) + + def test_gradient_flows(self) -> None: + """Tests that gradients propagate through the head.""" + head = self._make_head() + features = torch.randn(2, 16, self._EMBED_DIM, requires_grad=True) + logits = head(patch_features=features) + logits.sum().backward() + + self.assertIsNotNone(features.grad) + + def test_single_token(self) -> None: + """Tests forward pass with a single token.""" + head = self._make_head() + features = torch.randn(1, 1, self._EMBED_DIM) + logits = head(patch_features=features) + + self.assertEqual(logits.shape, (1, 1, self._VOCAB_SIZE)) + + @unittest.skipUnless(torch.cuda.is_available(), "CUDA not available") + def test_forward_pass_cuda(self) -> None: + """Tests forward pass on CUDA if available.""" + head = self._make_head().cuda() + features = torch.randn(2, 16, self._EMBED_DIM).cuda() + logits = head(patch_features=features) + + self.assertTrue(logits.is_cuda) + self.assertEqual(logits.shape, (2, 16, self._VOCAB_SIZE)) diff --git a/tests/transforms/test_beit_transform.py b/tests/transforms/test_beit_transform.py new file mode 100644 index 000000000..9fd8dc173 --- /dev/null +++ b/tests/transforms/test_beit_transform.py @@ -0,0 +1,188 @@ +from __future__ import annotations + +import unittest +from typing import Any + +import torch +from PIL import Image + +from lightly.transforms import BEiTTransform + + +class TestBEITTransform(unittest.TestCase): + """Tests for the BEiTTransform data augmentation and masking.""" + + _GRID = 14 + _N = 14 * 14 + _MASK_RATIO = 0.4 + + def _make_transform(self, **kwargs: Any) -> BEiTTransform: + """Creates a BEiTTransform with optional overrides. + + Args: + **kwargs: + Keyword arguments passed to BEiTTransform. + + Returns: + A BEiTTransform instance configured for testing. + """ + return BEiTTransform(**kwargs) + + def test_mask_shape(self) -> None: + """Tests that generated masks have the expected shape.""" + transform = self._make_transform() + for batch_size in [1, 4, 8]: + mask = transform.mask_generator(batch_size=batch_size) + self.assertEqual(mask.shape, (batch_size, self._N)) + + def test_mask_dtype_is_bool(self) -> None: + """Tests that generated masks have boolean dtype.""" + transform = self._make_transform() + mask = transform.mask_generator(batch_size=2) + self.assertEqual(mask.dtype, torch.bool) + + def test_mask_values_are_binary(self) -> None: + """Tests that mask values are strictly True or False.""" + transform = self._make_transform() + mask = transform.mask_generator(batch_size=4) + unique = mask.unique() + self.assertTrue(set(unique.tolist()).issubset({False, True})) + + def test_mask_is_non_empty(self) -> None: + """Tests that generated masks contain at least some masked patches.""" + transform = self._make_transform() + for _ in range(10): + mask = transform.mask_generator(batch_size=1) + self.assertGreater(mask[0].sum().item(), 0) + + def test_mask_does_not_cover_all_patches(self) -> None: + """Tests that generated masks do not cover all patches.""" + transform = self._make_transform() + for _ in range(10): + mask = transform.mask_generator(batch_size=1) + self.assertLess(mask[0].sum().item(), self._N) + + def test_masks_differ_across_samples_in_batch(self) -> None: + """Tests that masks differ across samples in the same batch.""" + transform = self._make_transform() + mask = transform.mask_generator(batch_size=4) + identical = all(torch.equal(mask[0], mask[i]) for i in range(1, 4)) + self.assertFalse(identical) + + def test_masks_differ_across_calls(self) -> None: + """Tests that masks differ across independent calls.""" + transform = self._make_transform() + mask_a = transform.mask_generator(batch_size=1) + mask_b = transform.mask_generator(batch_size=1) + self.assertFalse(torch.equal(mask_a, mask_b)) + + def test_custom_grid_size(self) -> None: + """Tests that custom input_size and patch_size produce correct grid.""" + transform = self._make_transform(input_size=112, patch_size=16) + mask = transform.mask_generator(batch_size=2) + self.assertEqual(mask.shape, (2, 49)) + + def test_custom_mask_ratio(self) -> None: + """Tests that custom mask_ratio is respected approximately.""" + ratio = 0.6 + transform = self._make_transform(mask_ratio=ratio) + target = int(ratio * self._N) + mask = transform.mask_generator(batch_size=4) + for i in range(4): + # Due to overlap-aware blocking, actual count may be slightly + # less than target, but should be close. + actual = mask[i].sum().item() + self.assertGreaterEqual(actual, int(target * 0.5)) + self.assertLessEqual(actual, self._N) + + def test_custom_min_block_patches(self) -> None: + """Tests that custom min_block_patches produces valid masks.""" + transform = self._make_transform(min_block_patches=1) + mask = transform.mask_generator(batch_size=4) + for i in range(4): + self.assertGreater(mask[i].sum().item(), 0) + + def test_max_block_patches(self) -> None: + """Tests that max_block_patches limits block size.""" + transform = self._make_transform(max_block_patches=4, min_block_patches=1) + mask = transform.mask_generator(batch_size=2) + for i in range(2): + self.assertGreater(mask[i].sum().item(), 0) + + def test_aspect_ratio_range(self) -> None: + """Tests that custom aspect_ratio_range is accepted.""" + transform = self._make_transform( + aspect_ratio_range=(0.5, 2.0), + ) + mask = transform.mask_generator(batch_size=2) + self.assertEqual(mask.shape, (2, self._N)) + + def test_mask_generator_repr(self) -> None: + """Tests that __repr__ returns a non-empty string.""" + transform = self._make_transform() + repr_str = repr(transform.mask_generator) + self.assertIsInstance(repr_str, str) + self.assertGreater(len(repr_str), 0) + + def test_mask_generator_get_shape(self) -> None: + """Tests that get_shape returns the expected grid dimensions.""" + transform = self._make_transform() + shape = transform.mask_generator.get_shape() + self.assertEqual(shape, (self._GRID, self._GRID)) + + def test_image_transform_output_shape(self) -> None: + """Tests that the image transform produces the correct tensor shape.""" + transform = self._make_transform() + image = Image.new(mode="RGB", size=(224, 224)) + tensor = transform(image=image) + self.assertEqual(tensor.shape, (3, 224, 224)) + + def test_image_transform_output_is_tensor(self) -> None: + """Tests that the image transform returns a torch.Tensor.""" + transform = self._make_transform() + image = Image.new(mode="RGB", size=(224, 224)) + tensor = transform(image=image) + self.assertIsInstance(tensor, torch.Tensor) + + def test_image_transform_normalized(self) -> None: + """Tests that the image transform produces normalized values.""" + transform = self._make_transform() + image = Image.new(mode="RGB", size=(224, 224)) + tensor = transform(image=image) + # After normalization, values should be outside [0, 1] for a black image + # or at least not exactly in [0, 1] range for typical images. + # For a black image (all zeros), normalized value = -mean/std. + self.assertFalse(torch.all((tensor >= 0) & (tensor <= 1))) + + def test_mask_generator_batch_consistency(self) -> None: + """Tests that batched generation is equivalent to individual calls.""" + transform = self._make_transform() + torch.manual_seed(42) + mask_batch = transform.mask_generator(batch_size=4) + + # Individual calls would differ due to randomness, so we just verify + # the batch dimension is correct. + self.assertEqual(mask_batch.shape[0], 4) + + def test_different_color_jitter(self) -> None: + """Tests that custom color_jitter is accepted.""" + transform = self._make_transform(color_jitter=0.2) + image = Image.new(mode="RGB", size=(224, 224)) + tensor = transform(image=image) + self.assertEqual(tensor.shape, (3, 224, 224)) + + def test_different_normalization_stats(self) -> None: + """Tests that custom mean and std are accepted.""" + transform = self._make_transform( + mean=(0.5, 0.5, 0.5), + std=(0.5, 0.5, 0.5), + ) + image = Image.new(mode="RGB", size=(224, 224)) + tensor = transform(image=image) + self.assertEqual(tensor.shape, (3, 224, 224)) + + def test_mask_generator_with_batch_size_one(self) -> None: + """Tests mask generation with the smallest batch size.""" + transform = self._make_transform() + mask = transform.mask_generator(batch_size=1) + self.assertEqual(mask.shape, (1, self._N))