{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n# Tutorial 8: Using timm Models as Backbones\n\nYou can use LightlySSL to pre-train any timm model using self-supervised learning since\nmost methods share a similar workflow. All methods have a base model (the backbone), which\ncan be any fundamental model such as ResNet or MobileNet.\n\nIn this tutorial, we will learn how to use a model architecture from the timm library\nas a backbone in a self-supervised learning workflow.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Imports\n\nImport the Python frameworks we need for this tutorial.\nMake sure you have the necessary packages installed.\n\n```console\npip install lightly\"[timm]\"\n```\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import timm\nimport torch\nimport torch.nn as nn"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "timm comes packaged with >700 pre-trained models designed to be flexible and easy to use.\nThese pre-trained models can be loaded using the\n[create_model()](https://huggingface.co/docs/timm/v1.0.8/en/reference/models#timm.create_model)\nfunction. For example, we can use the following snippet to create an efficient model.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "backbone = timm.create_model(model_name=\"efficientnet_b0\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Using a timm Model as a Backbone\n\nWe can now use this model as a backbone for training. Let's see how we can\ncreate a torch module for the [SimCLR](https://arxiv.org/abs/2002.05709) method.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "from lightly.models.modules.heads import SimCLRProjectionHead\n\n\nclass SimCLR(torch.nn.Module):\n    def __init__(self, backbone, feature_dim, out_dim):\n        super().__init__()\n        self.backbone = backbone\n        self.projection_head = SimCLRProjectionHead(feature_dim, feature_dim, out_dim)\n\n    def forward(self, x):\n        features = self.backbone.forward_features(x)\n        h = self.backbone.global_pool(features).flatten(start_dim=1)\n        z = self.projection_head(h)\n        return z\n\n\nsimclr = SimCLR(backbone, feature_dim=1280, out_dim=128)\n\n# check if it works\ninput_a = torch.randn((2, 3, 224, 224))\ninput_b = torch.randn((2, 3, 224, 224))\nout_a = simclr(input_a)\nout_b = simclr(input_b)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Next Steps\n\nFor an indepth tutorial on fine-tuning a model using [SimCLR](https://arxiv.org/abs/2002.05709)\nyou can refer to our fine-tuning `lightly-checkpoint-finetuning-tutorial-7`.\nInterested in pre-training your own self-supervised models? Check out our other\ntutorials:\n\n- `input-structure-label`\n- `lightly-moco-tutorial-2`\n- `lightly-simsiam-tutorial-4`\n- `lightly-custom-augmentation-5`\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
}