{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "\n.. role:: bash(code)\n :language: bash\n\n# Tutorial 6: Pre-train a Detectron2 Backbone with Lightly\n\nIn this tutorial we show how you can do self-supervised pre-training of a\n[Detectron2](https://github.com/facebookresearch/detectron2) backbone with lightly.\nThe focus of this tutorial is on how to get and store a pre-trained ResNet50\nbackbone of the popular Detectron2 framework. If you want to learn more about\nself-supervised learning in general, go check out the following tutorials:\n\n - `lightly-moco-tutorial-2`\n - `lightly-simclr-tutorial-3`:\n - `lightly-simsiam-tutorial-4`\n\nWhat you will learn:\n\n- How to retrieve a Detectron2 ResNet50 backbone for pre-training\n- How to do self-supervised learning with the backbone\n- How to store the backbone in a checkpoint file which can be used by Detectron2\n\n## Introduction\n\nFor many tasks in computer vision it can be beneficial to pre-train a neural network\non a domain-specific dataset prior to finetuning it. For example, a retail detection\nnetwork can be pre-trained with self-supervised learning on a large retail detection\ndataset. This way the neural network learns to extract relevant features from the images\nwithout requiring any annotations at all. As a consequence, it's possible to finetune\nthe network with only a handful of annotated images. This tutorial will guide you\nthrough the steps to pre-train a detection backbone from the popular\n[Detectron2](https://github.com/facebookresearch/detectron2) framework.\n\n## Prerequisites:\nFor the purpose of this tutorial you will need to:\n\n- Install Lightly: Follow the [instructions](https://docs.lightly.ai/getting_started/install.html)_.\n- Install Detectron2: Follow the [instructions](https://detectron2.readthedocs.io/en/latest/tutorials/install.html)_.\n- Download a dataset for pre-training (we will use the [Freiburg Groceries Dataset](https://github.com/PhilJd/freiburg_groceries_dataset) dataset). You can download it by cloning the Github repository and running `download_dataset.py`. Alternatively, you can use this [download link](http://aisdatasets.informatik.uni-freiburg.de/freiburg_groceries_dataset/freiburg_groceries_dataset.tar.gz)\n\n

Note

The [Freiburg Groceries Dataset](https://github.com/PhilJd/freiburg_groceries_dataset) consists of 5000 256x256 RGB images of 25 food classes.\n Images show one or multiple instances of grocery products in shelves or similar scenarios.

\n\nFinally, you will need the Detectron2 configuration files. They are available [here](https://github.com/facebookresearch/detectron2/tree/main/configs).\nIn this tutorial we use a Faster RCNN with a feature pyramid network (FPN), so make sure you have the relevant file (Base-RCNN-FPN.yaml) in your directory.\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Imports\n\nImport the Python frameworks we need for this tutorial.\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "import torch\nfrom detectron2 import config, modeling\nfrom detectron2.checkpoint import DetectionCheckpointer\n\nfrom lightly.data import LightlyDataset\nfrom lightly.loss import NTXentLoss\nfrom lightly.models.modules import SimCLRProjectionHead\nfrom lightly.transforms import SimCLRTransform" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Configuration\nLet's set the configuration parameters for our experiments.\n\nWe use a batch size of 512 and an input size of 128 in order to fit everything\non the available amount of memory on our GPU (16GB). The number of features\nis set to the default output size of the ResNet50 backbone.\n\nWe only train for 5 epochs because the focus of this tutorial is on the\nintegration with Detectron2.\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "num_workers = 8\nbatch_size = 512\ninput_size = 128\nnum_ftrs = 2048\n\nseed = 1\nmax_epochs = 5\n\n# use cuda if possible\ndevice = \"cuda\" if torch.cuda.is_available() else \"cpu\"" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "You might have downloaded the dataset somewhere else or are using a different one.\nSet the path to the dataset accordingly. Additionally, make sure to set the\npath to the config file of the Detectron2 model you want to use.\nWe will be using an RCNN with a feature pyramid network (FPN).\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "data_path = \"/datasets/freiburg_groceries_dataset/images\"\ncfg_path = \"./Base-RCNN-FPN.yaml\"" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Initialize the Detectron2 Model\n\nThe output of the Detectron2 ResNet50 backbone is a dictionary with the keys\n`res1` through `res5` (see the [documentation](https://detectron2.readthedocs.io/en/latest/modules/modeling.html#detectron2.modeling.ResNet)).\nThe keys correspond to the different stages of the ResNet. In this tutorial, we are only\ninterested in the high-level abstractions from the last layer, `res5`. Therefore,\nwe have to add an additional layer which picks the right output from the dictionary.\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "class SelectStage(torch.nn.Module):\n \"\"\"Selects features from a given stage.\"\"\"\n\n def __init__(self, stage: str = \"res5\"):\n super().__init__()\n self.stage = stage\n\n def forward(self, x):\n return x[self.stage]" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Let's load the config file and make some adjustments to ensure smooth training.\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "cfg = config.get_cfg()\ncfg.merge_from_file(cfg_path)\n\n# use cuda if possible\ncfg.MODEL.DEVICE = device\n\n# randomly initialize network\ncfg.MODEL.WEIGHTS = \"\"\n\n# detectron2 uses BGR by default but pytorch/torchvision use RGB\ncfg.INPUT.FORMAT = \"RGB\"" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Next, we can build the Detectron2 model and extract the ResNet50 backbone as\nfollows:\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "detmodel = modeling.build_model(cfg)\n\nsimclr_backbone = torch.nn.Sequential(\n detmodel.backbone.bottom_up,\n SelectStage(\"res5\"),\n # res5 has shape bsz x 2048 x 4 x 4\n torch.nn.AdaptiveAvgPool2d(1),\n).to(device)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "

Note

The Detectron2 ResNet is missing the average pooling layer used to get a tensor of shape bsz x 2048.\n Therefore, we add an average pooling as in the [PyTorch ResNet](https://github.com/pytorch/pytorch/blob/1022443168b5fad55bbd03d087abf574c9d2e9df/benchmarks/functional_autograd_benchmark/torchvision_models.py#L147).

\n\n\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Finally, let's build SimCLR around the backbone as shown in the other\ntutorials. For this, we only require an additional projection head.\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "projection_head = SimCLRProjectionHead(\n input_dim=num_ftrs,\n hidden_dim=num_ftrs,\n output_dim=128,\n).to(device)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Setup data augmentations and loaders\n\nWe start by defining the augmentations which should be used for training.\nWe use the same ones as in the SimCLR paper but change the input size and\nminimum scale of the random crop to adjust to our dataset.\n\nWe don't go into detail here about using the optimal augmentations.\nYou can learn more about the different augmentations and learned invariances\nhere: `lightly-advanced`.\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "transform = SimCLRTransform(input_size=input_size)\n\ndataset_train_simclr = LightlyDataset(input_dir=data_path, transform=transform)\n\ndataloader_train_simclr = torch.utils.data.DataLoader(\n dataset_train_simclr,\n batch_size=batch_size,\n shuffle=True,\n drop_last=True,\n num_workers=num_workers,\n)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Self-supervised pre-training\nNow all we need to do is define a loss and optimizer and start training!\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "criterion = NTXentLoss()\noptimizer = torch.optim.Adam(\n list(simclr_backbone.parameters()) + list(projection_head.parameters()),\n lr=1e-4,\n)\n\n\nfor e in range(max_epochs):\n mean_loss = 0.0\n for (x0, x1), _, _ in dataloader_train_simclr:\n x0 = x0.to(device)\n x1 = x1.to(device)\n\n y0 = projection_head(simclr_backbone(x0).flatten(start_dim=1))\n y1 = projection_head(simclr_backbone(x1).flatten(start_dim=1))\n\n # backpropagation\n loss = criterion(y0, y1)\n loss.backward()\n\n optimizer.step()\n optimizer.zero_grad()\n\n # update average loss\n mean_loss += loss.detach().cpu().item() / len(dataloader_train_simclr)\n\n print(f\"[Epoch {e:2d}] Mean Loss = {mean_loss:.2f}\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Storing the checkpoint\nNow, we can use the pre-trained backbone from the Detectron2 model. The code\nbelow shows how to save it as a Detectron2 checkpoint called `my_model.pth`.\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "# get the first module from the backbone (i.e. the detectron2 ResNet)\n# backbone:\n# L ResNet50\n# L SelectStage\n# L AdaptiveAvgPool2d\ndetmodel.backbone.bottom_up = simclr_backbone[0]\n\ncheckpointer = DetectionCheckpointer(detmodel, save_dir=\"./\")\ncheckpointer.save(\"my_model\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Finetuning with Detectron2\n\nThe checkpoint from above can now be used by any Detectron2 script. For example,\nyou can use the `train_net.py` script in the Detectron2 `tools`:\n\n\n\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "```none\npython train_net.py --config-file ../configs/COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_1x.yaml \\\n MODEL.WEIGHTS path/to/my_model.pth \\\n MODEL.PIXEL_MEAN 123.675,116.280,103.530 \\\n MODEL.PIXEL_STD 58.395,57.120,57.375 \\\n INPUT.FORMAT RGB\n```\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "The :py:class:`~lightly.transforms.simclr.SimCLRTransform` applies an ImageNet\nnormalization of the input images by default. Therefore, we have to normalize\nthe input images at training time, too. Since Detectron2 uses an input space\nin the range 0 - 255, we use the numbers above.\n\n\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "

Note

Since the model was pre-trained with images in the RGB input format, it's\n necessary to set the permute the order of the pixel mean, and pixel std as shown above.

\n\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Next Steps\n\n\nWant to learn more about our self-supervised models and how to choose\naugmentations properly? Check out our other tutorials:\n\n- `lightly-moco-tutorial-2`\n- `lightly-simclr-tutorial-3`\n- `lightly-simsiam-tutorial-4`\n- `lightly-detectron-tutorial-6`\n\n\n" ] } ], "metadata": { "kernelspec": { "display_name": "Python 3", "language": "python", "name": "python3" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.10.13" } }, "nbformat": 4, "nbformat_minor": 0 }