feat: add BEiT support

README.md

@@ -76,6 +76,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
 

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 <https://arxiv.org/abs/2106.08254>`_
+
+
+.. 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

docs/source/examples/models.rst

@@ -34,3 +34,4 @@ for PyTorch and PyTorch Lightning to give you a headstart when implementing your
     tico.rst
     vicreg.rst
     vicregl.rst
+    beit.rst

examples/notebooks/pytorch/beit.ipynb

@@ -0,0 +1,295 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "0",
+   "metadata": {},
+   "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."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "1",
+   "metadata": {
+    "lines_to_next_cell": 2
+   },
+   "source": [
+    "NOTE: The ImageTokenizer must be pre-trained before BEIT pre-training.\n",
+    "The tokenizer follows the DALL-E discrete VAE architecture and can be\n",
+    "trained as follows:\n",
+    "\n",
+    "    tokenizer = ImageTokenizer(vocab_size=8192)\n",
+    "    optimizer = torch.optim.Adam(tokenizer.parameters(), lr=1e-3)\n",
+    "\n",
+    "    for images in dataloader:\n",
+    "        logits, recon = tokenizer(images)\n",
+    "        loss = F.mse_loss(recon, images)  # Reconstruction loss\n",
+    "        loss.backward()\n",
+    "        optimizer.step()\n",
+    "        optimizer.zero_grad()\n",
+    "\n",
+    "    torch.save(tokenizer.state_dict(), \"tokenizer.pth\")\n",
+    "\n",
+    "During BEIT pre-training, load and freeze the tokenizer:\n",
+    "    tokenizer.load_state_dict(torch.load(\"tokenizer.pth\"))\n",
+    "    tokenizer.eval()\n",
+    "    for param in tokenizer.parameters():\n",
+    "        param.requires_grad = False"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "2",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from __future__ import annotations"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "3",
+   "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": "4",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from lightly.loss import MaskedImageModelingLoss\n",
+    "from lightly.models.modules import BEITEncoder, ImageTokenizer, MIMHead\n",
+    "from lightly.transforms import BEITTransform"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "5",
+   "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: ImageTokenizer,\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 = MIMHead(\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": "6",
+   "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 = ImageTokenizer(vocab_size=512)\n",
+    "model = BEIT(encoder=encoder, tokenizer=tokenizer)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "7",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
+    "model.to(device=device)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "8",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "transform = BEITTransform(input_size=224, patch_size=8)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "9",
+   "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": "10",
+   "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": "11",
+   "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": "12",
+   "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": "13",
+   "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
+}