SparK Implementation

docs/source/lightly.loss.rst

@@ -56,6 +56,9 @@ lightly.loss
 .. autoclass:: lightly.loss.regularizer.co2.CO2Regularizer
    :members:
 
+.. autoclass:: lightly.loss.sparse_spark.SparKPatchReconLoss
+   :members:
+
 .. autoclass:: lightly.loss.swav_loss.SwaVLoss
    :members:
 

examples/notebooks/pytorch_lightning/spark.ipynb

@@ -0,0 +1,302 @@
+{
+ "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 timm\n",
+    "import torch\n",
+    "import torchvision\n",
+    "from pytorch_lightning import LightningModule\n",
+    "from torch import Tensor\n",
+    "from torch.nn import Module\n",
+    "from torchvision.transforms import v2"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "4",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import lightly.models.utils as model_utils\n",
+    "from lightly.loss.sparse_spark import SparKPatchReconLoss\n",
+    "from lightly.models.modules import sparse_spark"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "5",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "## The global projection head is the same as the Barlow Twins one\n",
+    "from lightly.models.modules.sparse_spark import (\n",
+    "    SparKDensifier,\n",
+    "    SparKMasker,\n",
+    "    SparKMaskingOutput,\n",
+    "    SparKOutputDecoder,\n",
+    "    UNetDecoder,\n",
+    "    sparse_layer_context,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "6",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def _get_downsample_ratio_from_timm_model(model: Module) -> int:\n",
+    "    if not hasattr(model, \"feature_info\"):\n",
+    "        raise ValueError(\n",
+    "            \"The provided model does not have the required 'feature_info' attribute.\"\n",
+    "        )\n",
+    "    return model.feature_info[-1][\"reduction\"]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "7",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def _get_enc_feat_map_chs_from_timm_model(model: Module) -> list[int]:\n",
+    "    if not hasattr(model, \"feature_info\"):\n",
+    "        raise ValueError(\n",
+    "            \"The provided model does not have the required 'feature_info' attribute.\"\n",
+    "        )\n",
+    "    return [fi[\"num_chs\"] for fi in model.feature_info]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "8",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "class SparseSparK(LightningModule):\n",
+    "    def __init__(\n",
+    "        self,\n",
+    "        input_size: int = 416,\n",
+    "        mask_ratio: float = 0.6,\n",
+    "        densify_norm: str = \"bn\",\n",
+    "        sbn=False,\n",
+    "    ):\n",
+    "        super().__init__()\n",
+    "        backbone = timm.create_model(\n",
+    "            model_name=\"resnet18\", drop_path_rate=0.05, features_only=True\n",
+    "        )\n",
+    "        downsample_ratio = _get_downsample_ratio_from_timm_model(backbone)\n",
+    "        enc_feat_map_chs = _get_enc_feat_map_chs_from_timm_model(backbone)\n",
+    "        self.sparse_encoder = sparse_spark.dense_model_to_sparse(\n",
+    "            m=backbone, sbn=sbn, verbose=True\n",
+    "        )\n",
+    "        self.fmap_h = input_size // downsample_ratio\n",
+    "        self.fmap_w = input_size // downsample_ratio\n",
+    "        self.dense_decoder = UNetDecoder(\n",
+    "            up_sample_ratio=downsample_ratio,\n",
+    "            width=enc_feat_map_chs[-1],\n",
+    "        )\n",
+    "        self.masker = SparKMasker(\n",
+    "            feature_map_size=(self.fmap_h, self.fmap_w),\n",
+    "            downsample_ratio=downsample_ratio,\n",
+    "            mask_ratio=mask_ratio,\n",
+    "        )\n",
+    "        self.densifier = SparKDensifier(\n",
+    "            encoder_in_channels=enc_feat_map_chs,\n",
+    "            decoder_in_channel=self.dense_decoder.width,\n",
+    "            densify_norm_str=densify_norm.lower(),\n",
+    "            sbn=sbn,\n",
+    "        )\n",
+    "        self.downsample_ratio = downsample_ratio\n",
+    "        # loss module for patch reconstruction\n",
+    "        self.recon_loss_fn = SparKPatchReconLoss()\n",
+    "        # output decoder for visualization (pass minimal spatial props)\n",
+    "        self.output_decoder = SparKOutputDecoder(\n",
+    "            fmap_h=self.fmap_h,\n",
+    "            fmap_w=self.fmap_w,\n",
+    "            downsample_ratio=downsample_ratio,\n",
+    "        )\n",
+    "\n",
+    "    def forward(\n",
+    "        self,\n",
+    "        inp_bchw: Tensor,\n",
+    "        vis=False,\n",
+    "    ):\n",
+    "        # step1. Mask\n",
+    "        mask_out: SparKMaskingOutput = self.masker(inp_bchw)\n",
+    "        masked_bchw, per_level_mask = mask_out\n",
+    "        active_b1fHfW = per_level_mask[0]\n",
+    "        active_b1hw = per_level_mask[-1]\n",
+    "        # step2. Encode: get hierarchical encoded sparse features (a list containing 4 feature maps at 4 scales)\n",
+    "        # Use sparse_layer_context to provide the mask to the sparse encoder and densifier.\n",
+    "        with sparse_layer_context(active_mask=active_b1fHfW):\n",
+    "            fea_bcffs: list[Tensor] = self.sparse_encoder(masked_bchw)\n",
+    "            # step3. Densify: get hierarchical dense features for decoding\n",
+    "            to_dec = self.densifier(fea_bcffs)\n",
+    "        # step4. Decode and reconstruct\n",
+    "        rec_bchw = self.dense_decoder(to_dec)\n",
+    "        inp, rec = (\n",
+    "            model_utils.patchify(inp_bchw, self.downsample_ratio),\n",
+    "            model_utils.patchify(rec_bchw, self.downsample_ratio),\n",
+    "        )  # inp and rec: (B, L = f*f, N = C*downsample_ratio**2)\n",
+    "\n",
+    "        recon_loss, mean, var = self.recon_loss_fn(\n",
+    "            inp_patches=inp, rec_patches=rec, active_mask=active_b1fHfW\n",
+    "        )\n",
+    "\n",
+    "        if vis:\n",
+    "            return self.output_decoder(\n",
+    "                rec_patches=rec,\n",
+    "                mean=mean,\n",
+    "                var=var,\n",
+    "                inp_bchw=inp_bchw,\n",
+    "                active_mask_full=active_b1hw,\n",
+    "            )\n",
+    "        else:\n",
+    "            return recon_loss\n",
+    "\n",
+    "    def training_step(self, batch: tuple[Tensor, Tensor], batch_index: int) -> Tensor:\n",
+    "        img, _ = batch\n",
+    "        recon_loss = self.forward(img)\n",
+    "        # Log the training loss to logger and progress bar (per-step and per-epoch)\n",
+    "        self.log(\n",
+    "            \"train_loss\",\n",
+    "            recon_loss,\n",
+    "            on_step=True,\n",
+    "            on_epoch=True,\n",
+    "            prog_bar=True,\n",
+    "            logger=True,\n",
+    "        )\n",
+    "        return recon_loss\n",
+    "\n",
+    "    def configure_optimizers(self):\n",
+    "        return torch.optim.SGD(\n",
+    "            self.parameters(), lr=0.03, momentum=0.9, weight_decay=1e-4\n",
+    "        )"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "9",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "model = SparseSparK(input_size=416)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "10",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "dataset = torchvision.datasets.Caltech101(\n",
+    "    \"datasets/caltech101\",\n",
+    "    download=True,\n",
+    "    transform=v2.Compose(\n",
+    "        [\n",
+    "            v2.Resize((416, 416)),\n",
+    "            v2.RGB(),\n",
+    "            v2.ToTensor(),\n",
+    "            v2.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),\n",
+    "        ]\n",
+    "    ),\n",
+    ")\n",
+    "# or create a dataset from a folder containing images or videos:\n",
+    "# dataset = LightlyDataset(\"path/to/folder\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "11",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "dataloader = torch.utils.data.DataLoader(\n",
+    "    dataset,\n",
+    "    batch_size=4,\n",
+    "    shuffle=True,\n",
+    "    drop_last=True,\n",
+    "    num_workers=8,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "12",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "trainer = pl.Trainer(\n",
+    "    max_epochs=30,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "13",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "trainer.fit(\n",
+    "    model=model,\n",
+    "    train_dataloaders=dataloader,\n",
+    ")"
+   ]
+  }
+ ],
+ "metadata": {
+  "jupytext": {
+   "cell_metadata_filter": "-all",
+   "main_language": "python",
+   "notebook_metadata_filter": "-all"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}