Repository: lightly-ai/lightly

Issue:
pretraining Detectron2 FPN models
Hi Lightly-team,

Is there a way to use Lightly for self-supervised pretraining of Detectron2 models?
I think that would be valuable for many people.
The "Lightly Backbone Playground.ipynb" colab shows the concept of using a Lightly implemented SimCLR to pretrain a torchvision resnet50.
However, the most effective Detectron2 models are resnet-FPNs (Feature Pyramid Networks).

The detectron standard Faster RCNN with resnet-FPN contains a backbone: detmodel.backbone
a RPN (region proposal network): detmodel.proposal_generator
and a ROI (region of interest): detmodel.head roi_heads
The backbone itself is a resnet-FPN and I assume it is easiest to 
train the bottom-up part of the resnet-FPN ( detmodel.backbone.bottom_up), as that is basically a resnet.

It seems doable except that the feature-shapes are mismatched.
I created a minimal working example in colab:
https://colab.research.google.com/drive/11AIdCGkHC1loQjvQnj51kXCD9TKNAib5?usp=sharing

Is there a way to get the SimCLR working with this backbone?

The colab above contains a runnable example, but I tried to capture the essence in the code below:
```
# Get original resnet model
resnet50 = torchvision.models.resnet50()
# Remove linear head
rbackbone = nn.Sequential(*list(resnet50.children())[:-1],)

# Toy input
toy_input_view_a = torch.randn((11,3,256,256)).to(device)

# Output shape of original resnet model
np.shape( rbackbone(toy_input_view_a) )
# Out: torch.Size([11, 2048, 1, 1])

# In the forward function of the file lightly/models/simclr.py:88
# f0 = self.backbone(x0).flatten(start_dim=1)
# which leads to the shape torch.Size([11, 2048]) for rbackbone's output

#However, with the detectron model
detmodel = build_model(<config_file>)
backbone = nn.Sequential(*list(detmodel.backbone.bottom_up.children()),)

# We get a different output shape to start with
np.shape( backbone(toy_input_view_a) )
# Out: torch.Size([11, 2048, 8, 8])

# Which, after the flatten in lightly simclr model lightly/models/simclr.py:88
# leads to the shape torch.Size([11, 131072])

# Thus 
simCLR = lightly.models.SimCLR(backbone, num_ftrs=2048)
out = simCLR(toy_input_view_a, toy_input_view_a)
# Leads to the error
# RuntimeError: mat1 and mat2 shapes cannot be multiplied (11x131072 and 2048x2048)
```


Discussion:
Hi @RafaelMostert

Pre-training the Detectron2 backbone with Lightly is certainly possible. In general, any backbone can be used with only slight modifications (could be worth adding helper functions for this @IgorSusmelj and @guarin).

Regarding your problem: When you specify `SimCLR` with `lightly.models.SimCLR(backbone, num_ftrs=2048)` the projection head in `SimCLR` expects a 2048-dimensional output from the backbone. However, with flattening, the output from your backbone is 131072-dimensional. What we typically do to prevent this is to add a [`AdaptiveAvgPool2d`](https://pytorch.org/docs/stable/generated/torch.nn.AdaptiveAvgPool2d.html) layer after the backbone. This will transform the `[11, 2048, 8, 8]` tensor into a `[11, 2048]` tensor.
[Here](https://docs.lightly.ai/tutorials/package/tutorial_moco_memory_bank.html#create-the-moco-lightning-module)'s an example. In the code for creating the backbone in the `MoCoModel`, we use

```python
backbone = nn.Sequential(
    *list(resnet.children())[:-1], # output is [bsz, 2048, 8, 8]
    nn.AdaptiveAvgPool2d(1), # output is [bsz, 2048]
)
```


Thanks, that works!

There is enough documentation for the training part.

The following step after training would be loading the simCLR pre-trained weights back into the Detectron2 model.
For people dealing with something similar. The code below does the job for me.
The small difference with the Lightly tutorials lies in changing the key-names of the pytorch state_dict and using `load_state_dict` with the `strict=False` parameter.

```
# Storing the simCLR pretrained backbone without the added averagepool-layer
state_dict = {'resnet50_parameters': simCLR.backbone[:-1].state_dict()}
torch.save(state_dict, 'pretrained_model.pth')

######### Usually in a new file ####################
# Initialize a detectron2 model
new_detmodel = build_model(cfg)
# Load the pretrained model
ckpt = torch.load('pretrained_model.pth')
state = ckpt['resnet50_parameters']

# Change keynames of simCLR pretrained model state_dict to match detectron2 state_dict
# Sidenote: Iterating over state is equivalent to iterating over state.keys()
pretrained_keys = deepcopy(list(state.keys()))
destination_keys = deepcopy(list(new_detmodel.backbone.bottom_up.state_dict()))
for i, (old_key, dest_key) in enumerate(zip(pretrained_keys, destination_keys)): 
    if i<3: print(old_key,' ,  ',dest_key) # Show the slight difference in keynames
    assert old_key.split('.')[1:] == dest_key.split('.')[1:]
    state[dest_key] = state[old_key]
[state.pop(k) for k in pretrained_keys]; # Delete old keys

# Load partial model weights
new_detmodel.backbone.bottom_up.load_state_dict(state,strict=False)
```
Finally, I wonder, would it be possible (and sensible?) to pre-train the whole Feature Pyramid Network backbone with Lightly instead of just the ResNet (bottom_up) within it?

Below is a view of the architecture
```
FPN(
  (fpn_lateral2): Conv2d(256, 256, kernel_size=(1, 1), stride=(1, 1))
  (fpn_output2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  (fpn_lateral3): Conv2d(512, 256, kernel_size=(1, 1), stride=(1, 1))
  (fpn_output3): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  (fpn_lateral4): Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1))
  (fpn_output4): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  (fpn_lateral5): Conv2d(2048, 256, kernel_size=(1, 1), stride=(1, 1))
  (fpn_output5): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
  (top_block): LastLevelMaxPool()
  (bottom_up): ResNet(.....etc) # Training this part is already discussed in previous comments  
)
```
Using the forward pass of the backbone returns a dict with the feature maps "p2,p3,p4,p5,p6" 
as detailed in section 3 of https://arxiv.org/pdf/1612.03144.pdf

```
result = detmodel.backbone(toy_input_view_a)

for k in result.keys():
    print(f'key: {k}, shape {np.shape(result[k])}')
# out: key p2, shape torch.Size([11, 256, 64, 64])
# out: key p3, shape torch.Size([11, 256, 32, 32])
# out: key p4, shape torch.Size([11, 256, 16, 16])
# out: key p5, shape torch.Size([11, 256, 8, 8])
# out: key p6, shape torch.Size([11, 256, 4, 4])

```
Which, when naively inserted into the simCLR, would lead to an error at the flatten function in lightly/models/simclr.py:88
If you're interested in experimenting you could simply concatenate the outputs into a single output to prevent the error in `SimCLR`. However, I doubt that training the full FPN would yield meaningful results.

There is a lot of ongoing research about pre-training object detection models in a self-supervised manner and I'm happy to point you in that direction if that's something you want to look into.
> I doubt that training the full FPN would yield meaningful results.

It is indeed a very small extra part of the backbone, so it is probably not of much added value.

> There is a lot of ongoing research about pre-training object detection models in a self-supervised manner and I'm happy to point you in that direction if that's something you want to look into.

That would be great.




We are working on a new tutorial to address this. Hopefully, this will be out next week!

- How to pretrain a detectron2 model using Lightly
- How to use the pretrained model to fine-tune a detectron2 model

Resolving patch:
diff --git a/docs/source/tutorials_source/package/tutorial_pretrain_detectron2.py b/docs/source/tutorials_source/package/tutorial_pretrain_detectron2.py
new file mode 100644
--- /dev/null
+++ b/docs/source/tutorials_source/package/tutorial_pretrain_detectron2.py
@@ -0,0 +1,284 @@
+"""
+.. _lightly-detectron-tutorial-6:
+
+.. role:: bash(code)
+   :language: bash
+
+Tutorial 6: Pre-train a Detectron2 Backbone with Lightly
+==========================================================
+
+In this tutorial we show how you can do self-supervised pre-training of a
+`Detectron2 <https://github.com/facebookresearch/detectron2>`_ backbone with lightly.
+The focus of this tutorial is on how to get and store a pre-trained ResNet50
+backbone of the popular Detectron2 framework. If you want to learn more about
+self-supervised learning in general, go check out the following tutorials:
+
+ - :ref:`lightly-moco-tutorial-2`
+ - :ref:`lightly-simclr-tutorial-3`:
+ - :ref:`lightly-simsiam-tutorial-4`
+
+What you will learn:
+
+- How to retrieve a Detectron2 ResNet50 backbone for pre-training
+- How to do self-supervised learning with the backbone
+- How to store the backbone in a checkpoint file which can be used by Detectron2
+
+Introduction
+----------------
+
+For many tasks in computer vision it can be beneficial to pre-train a neural network
+on a domain-specific dataset prior to finetuning it. For example, a retail detection
+network can be pre-trained with self-supervised learning on a large retail detection
+dataset. This way the neural network learns to extract relevant features from the images
+without requiring any annotations at all. As a consequence, it's possible to finetune
+the network with only a handful of annotated images. This tutorial will guide you
+through the steps to pre-train a detection backbone from the popular
+`Detectron2 <https://github.com/facebookresearch/detectron2>`_ framework.
+
+Prerequisites:
+----------------
+For the purpose of this tutorial you will need to:
+
+- Install Lightly: Follow the `instructions <https://docs.lightly.ai/getting_started/install.html>`_.
+- Install Detectron2: Follow the `instructions <https://detectron2.readthedocs.io/en/latest/tutorials/install.html>`_.
+- 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>`_
+
+.. note::
+
+    The `Freiburg Groceries Dataset <https://github.com/PhilJd/freiburg_groceries_dataset>`_ consists of 5000 256x256 RGB images of 25 food classes.
+    Images show one or multiple instances of grocery products in shelves or similar scenarios.
+
+Finally, you will need the Detectron2 configuration files. They are available `here <https://github.com/facebookresearch/detectron2/tree/main/configs>`_.
+In 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.
+
+
+"""
+
+# %%
+# Imports
+# -------
+#
+# Import the Python frameworks we need for this tutorial.
+import torch
+import lightly
+from detectron2 import config, modeling
+from detectron2.checkpoint import DetectionCheckpointer
+
+# %%
+# Configuration
+# -------------
+# Let's set the configuration parameters for our experiments.
+#
+# We use a batch size of 512 and an input size of 128 in order to fit everything
+# on the available amount of memory on our GPU (16GB). The number of features
+# is set to the default output size of the ResNet50 backbone.
+#
+# We only train for 5 epochs because the focus of this tutorial is on the
+# integration with Detectron2.
+
+num_workers = 8
+batch_size = 512
+input_size = 128
+num_ftrs = 2048
+
+seed = 1
+max_epochs = 0
+
+# use cuda if possible
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+
+
+# %%
+# You might have downloaded the dataset somewhere else or are using a different one.
+# Set the path to the dataset accordingly. Additionally, make sure to set the
+# path to the config file of the Detectron2 model you want to use.
+# We will be using an RCNN with a feature pyramid network (FPN).
+data_path = '/datasets/freiburg_groceries_dataset/images'
+cfg_path = './Base-RCNN-FPN.yaml'
+
+# %%
+# Initialize the Detectron2 Model
+# --------------------------------
+# 
+# The output of the Detectron2 ResNet50 backbone is a dictionary with the keys
+# `res1` through `res5` (see the `documentation <https://detectron2.readthedocs.io/en/latest/modules/modeling.html#detectron2.modeling.ResNet>`_).
+# The keys correspond to the different stages of the ResNet. In this tutorial, we are only
+# interested in the high-level abstractions from the last layer, `res5`. Therefore,
+# we have to add an additional layer which picks the right output from the dictionary.
+class SelectStage(torch.nn.Module):
+    """Selects features from a given stage."""
+
+    def __init__(self, stage: str = 'res5'):
+        super().__init__()
+        self.stage = stage
+
+    def forward(self, x):
+        return x[self.stage]
+
+# %%
+# Let's load the config file and make some adjustments to ensure smooth training.
+cfg = config.get_cfg()
+cfg.merge_from_file(cfg_path)
+
+# use cuda if possible
+cfg.MODEL.DEVICE = device
+
+# randomly initialize network
+cfg.MODEL.WEIGHTS = ""
+
+# detectron2 uses BGR by default but pytorch/torchvision use RGB
+cfg.INPUT.FORMAT = "RGB"
+
+# %%
+# Next, we can build the Detectron2 model and extract the ResNet50 backbone as
+# follows:
+
+detmodel = modeling.build_model(cfg)
+
+simclr_backbone = torch.nn.Sequential(
+    detmodel.backbone.bottom_up,
+    SelectStage('res5'),
+    # res5 has shape bsz x 2048 x 4 x 4
+    torch.nn.AdaptiveAvgPool2d(1),
+).to(device)
+
+# %%
+#
+#
+#.. note::
+#
+#   The Detectron2 ResNet is missing the average pooling layer used to get a tensor of shape bsz x 2048.
+#   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>`_.
+#
+
+# %%
+# Finally, let's build SimCLR around the backbone as shown in the other
+# tutorials. For this, we only require an additional projection head.
+projection_head = lightly.models.modules.SimCLRProjectionHead(
+    input_dim=num_ftrs,
+    hidden_dim=num_ftrs,
+    output_dim=128,
+).to(device)
+
+# %%
+# Setup data augmentations and loaders
+# ------------------------------------
+#
+# We start by defining the augmentations which should be used for training.
+# We use the same ones as in the SimCLR paper but change the input size and
+# minimum scale of the random crop to adjust to our dataset. 
+#
+# We don't go into detail here about using the optimal augmentations.
+# You can learn more about the different augmentations and learned invariances
+# here: :ref:`lightly-advanced`.
+collate_fn = lightly.data.SimCLRCollateFunction(input_size=input_size)
+
+dataset_train_simclr = lightly.data.LightlyDataset(input_dir=data_path)
+
+dataloader_train_simclr = torch.utils.data.DataLoader(
+    dataset_train_simclr,
+    batch_size=batch_size,
+    shuffle=True,
+    collate_fn=collate_fn,
+    drop_last=True,
+    num_workers=num_workers
+)
+
+# %%
+# Self-supervised pre-training
+# -----------------------------
+# Now all we need to do is define a loss and optimizer and start training!
+
+criterion = lightly.loss.NTXentLoss()
+optimizer = torch.optim.Adam(
+    list(simclr_backbone.parameters()) + list(projection_head.parameters()),
+    lr=1e-4,
+)
+
+
+for e in range(max_epochs):
+
+    mean_loss = 0.
+    for (x0, x1), _, _ in dataloader_train_simclr:
+
+        x0 = x0.to(device)
+        x1 = x1.to(device)
+
+        y0 = projection_head(simclr_backbone(x0).flatten(start_dim=1))
+        y1 = projection_head(simclr_backbone(x1).flatten(start_dim=1))
+
+        # backpropagation
+        loss = criterion(y0, y1)
+        loss.backward()
+
+        optimizer.step()
+        optimizer.zero_grad()
+
+        # update average loss
+        mean_loss += loss.detach().cpu().item() / len(dataloader_train_simclr)
+
+    print(f'[Epoch {e:2d}] Mean Loss = {mean_loss:.2f}')
+
+
+# %%
+# Storing the checkpoint
+# -----------------------
+# Now, we can use the pre-trained backbone from the Detectron2 model. The code
+# below shows how to save it as a Detectron2 checkpoint called `my_model.pth`.
+
+# get the first module from the backbone (i.e. the detectron2 ResNet)
+# backbone:
+#     L ResNet50
+#     L SelectStage
+#     L AdaptiveAvgPool2d
+detmodel.backbone.bottom_up = simclr_backbone[0]
+
+checkpointer = DetectionCheckpointer(detmodel, save_dir='./')
+checkpointer.save('my_model')
+
+
+# %%
+# Finetuning with Detectron2
+# ---------------------------
+#
+# The checkpoint from above can now be used by any Detectron2 script. For example,
+# you can use the `train_net.py` script in the Detectron2 `tools`:
+#
+#
+
+# %%
+#.. code-block:: none
+#
+#   python train_net.py --config-file ../configs/COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_1x.yaml \
+#       MODEL.WEIGHTS path/to/my_model.pth \
+#       MODEL.PIXEL_MEAN 123.675,116.280,103.530 \
+#       MODEL.PIXEL_STD 58.395,57.120,57.375 \
+#       INPUT.FORMAT RGB
+#
+
+# %%
+# 
+# The :py:class:`lightly.data.collate.SimCLRCollateFunction` applies an ImageNet
+# normalization of the input images by default. Therefore, we have to normalize
+# the input images at training time, too. Since Detectron2 uses an input space
+# in the range 0 - 255, we use the numbers above.
+# 
+
+# %%
+#
+#.. note::
+#
+#   Since the model was pre-trained with images in the RGB input format, it's
+#   necessary to set the permute the order of the pixel mean, and pixel std as shown above.
+
+# %%
+# Next Steps
+# ------------
+#
+#
+# Want to learn more about our self-supervised models and how to choose
+# augmentations properly? Check out our other tutorials:
+#
+# - :ref:`lightly-moco-tutorial-2`
+# - :ref:`lightly-simclr-tutorial-3`
+# - :ref:`lightly-simsiam-tutorial-4`


Test patch: