{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n# Tutorial 4: Train SimSiam on Satellite Images\n\nIn this tutorial we will train a SimSiam model in old-school PyTorch style on a\nset of satellite images of Italy. We will showcase how the generated embeddings\ncan be used for exploration and better understanding of the raw data.\n\nYou can read up on the model in the paper\n[Exploring Simple Siamese Representation Learning](https://arxiv.org/abs/2011.10566).\n\nWe will be using a dataset of satellite images from ESAs Sentinel-2 satellite over Italy.\nIf you're interested, you can get your own data from the [Copernicus Open Access Hub](https://scihub.copernicus.eu/).\nThe original images have been cropped into smaller tiles due to their immense size and\nthe dataset has been balanced based on a simple clustering of the mean RGB color values\nto prevent a surplus of images of the sea.\n\nIn this tutorial you will learn:\n\n- How to work with the SimSiam model\n\n- How to do self-supervised learning using PyTorch\n\n- How to check whether your embeddings have collapsed\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 math\n\nimport numpy as np\nimport torch\nimport torch.nn as nn\nimport torchvision\n\nfrom lightly.data import LightlyDataset\nfrom lightly.loss import NegativeCosineSimilarity\nfrom lightly.models.modules.heads import SimSiamPredictionHead, SimSiamProjectionHead\nfrom lightly.transforms import SimCLRTransform, utils"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Configuration\n\nWe set some configuration parameters for our experiment.\n\nThe default configuration with a batch size and input resolution of 256\nrequires 16GB of GPU memory.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "num_workers = 8\nbatch_size = 128\nseed = 1\nepochs = 50\ninput_size = 256\n\n# dimension of the embeddings\nnum_ftrs = 512\n# dimension of the output of the prediction and projection heads\nout_dim = proj_hidden_dim = 512\n# the prediction head uses a bottleneck architecture\npred_hidden_dim = 128"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Let's set the seed for our experiments and the path to our data\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# seed torch and numpy\ntorch.manual_seed(0)\nnp.random.seed(0)\n\n# set the path to the dataset\npath_to_data = \"/datasets/sentinel-2-italy-v1/\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Setup data augmentations and loaders\nSince we're working on satellite images, it makes sense to use horizontal and\nvertical flips as well as random rotation transformations. We apply weak color\njitter to learn an invariance of the model with respect to slight changes in\nthe color of the water.\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# define the augmentations for self-supervised learning\ntransform = SimCLRTransform(\n    input_size=input_size,\n    # require invariance to flips and rotations\n    hf_prob=0.5,\n    vf_prob=0.5,\n    rr_prob=0.5,\n    # satellite images are all taken from the same height\n    # so we use only slight random cropping\n    min_scale=0.5,\n    # use a weak color jitter for invariance w.r.t small color changes\n    cj_prob=0.2,\n    cj_bright=0.1,\n    cj_contrast=0.1,\n    cj_hue=0.1,\n    cj_sat=0.1,\n)\n\n# create a lightly dataset for training with augmentations\ndataset_train_simsiam = LightlyDataset(input_dir=path_to_data, transform=transform)\n\n# create a dataloader for training\ndataloader_train_simsiam = torch.utils.data.DataLoader(\n    dataset_train_simsiam,\n    batch_size=batch_size,\n    shuffle=True,\n    drop_last=True,\n    num_workers=num_workers,\n)\n\n# create a torchvision transformation for embedding the dataset after training\n# here, we resize the images to match the input size during training and apply\n# a normalization of the color channel based on statistics from imagenet\ntest_transforms = torchvision.transforms.Compose(\n    [\n        torchvision.transforms.Resize((input_size, input_size)),\n        torchvision.transforms.ToTensor(),\n        torchvision.transforms.Normalize(\n            mean=utils.IMAGENET_NORMALIZE[\"mean\"],\n            std=utils.IMAGENET_NORMALIZE[\"std\"],\n        ),\n    ]\n)\n\n# create a lightly dataset for embedding\ndataset_test = LightlyDataset(input_dir=path_to_data, transform=test_transforms)\n\n# create a dataloader for embedding\ndataloader_test = torch.utils.data.DataLoader(\n    dataset_test,\n    batch_size=batch_size,\n    shuffle=False,\n    drop_last=False,\n    num_workers=num_workers,\n)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Create the SimSiam model\n\nCreate a ResNet backbone and remove the classification head\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "class SimSiam(nn.Module):\n    def __init__(self, backbone, num_ftrs, proj_hidden_dim, pred_hidden_dim, out_dim):\n        super().__init__()\n        self.backbone = backbone\n        self.projection_head = SimSiamProjectionHead(num_ftrs, proj_hidden_dim, out_dim)\n        self.prediction_head = SimSiamPredictionHead(out_dim, pred_hidden_dim, out_dim)\n\n    def forward(self, x):\n        # get representations\n        f = self.backbone(x).flatten(start_dim=1)\n        # get projections\n        z = self.projection_head(f)\n        # get predictions\n        p = self.prediction_head(z)\n        # stop gradient\n        z = z.detach()\n        return z, p\n\n\n# we use a pretrained resnet for this tutorial to speed\n# up training time but you can also train one from scratch\nresnet = torchvision.models.resnet18()\nbackbone = nn.Sequential(*list(resnet.children())[:-1])\nmodel = SimSiam(backbone, num_ftrs, proj_hidden_dim, pred_hidden_dim, out_dim)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "SimSiam uses a symmetric negative cosine similarity loss and does therefore\nnot require any negative samples. We build a criterion and an optimizer.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# SimSiam uses a symmetric negative cosine similarity loss\ncriterion = NegativeCosineSimilarity()\n\n# scale the learning rate\nlr = 0.05 * batch_size / 256\n# use SGD with momentum and weight decay\noptimizer = torch.optim.SGD(model.parameters(), lr=lr, momentum=0.9, weight_decay=5e-4)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Train SimSiam\n\nTo train the SimSiam model, you can use a classic PyTorch training loop:\nFor every epoch, iterate over all batches in the training data, extract\nthe two transforms of every image, pass them through the model, and calculate\nthe loss. Then, simply update the weights with the optimizer. Don't forget to\nreset the gradients!\n\nSince SimSiam doesn't require negative samples, it is a good idea to check\nwhether the outputs of the model have collapsed into a single direction. For\nthis we can simply check the standard deviation of the L2 normalized output\nvectors. If it is close to one divided by the square root of the output\ndimension, everything is fine (you can read\nup on this idea [here](https://arxiv.org/abs/2011.10566)).\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\nmodel.to(device)\n\navg_loss = 0.0\navg_output_std = 0.0\nfor e in range(epochs):\n    for (x0, x1), _, _ in dataloader_train_simsiam:\n        # move images to the gpu\n        x0 = x0.to(device)\n        x1 = x1.to(device)\n\n        # run the model on both transforms of the images\n        # we get projections (z0 and z1) and\n        # predictions (p0 and p1) as output\n        z0, p0 = model(x0)\n        z1, p1 = model(x1)\n\n        # apply the symmetric negative cosine similarity\n        # and run backpropagation\n        loss = 0.5 * (criterion(z0, p1) + criterion(z1, p0))\n        loss.backward()\n\n        optimizer.step()\n        optimizer.zero_grad()\n\n        # calculate the per-dimension standard deviation of the outputs\n        # we can use this later to check whether the embeddings are collapsing\n        output = p0.detach()\n        output = torch.nn.functional.normalize(output, dim=1)\n\n        output_std = torch.std(output, 0)\n        output_std = output_std.mean()\n\n        # use moving averages to track the loss and standard deviation\n        w = 0.9\n        avg_loss = w * avg_loss + (1 - w) * loss.item()\n        avg_output_std = w * avg_output_std + (1 - w) * output_std.item()\n\n    # the level of collapse is large if the standard deviation of the l2\n    # normalized output is much smaller than 1 / sqrt(dim)\n    collapse_level = max(0.0, 1 - math.sqrt(out_dim) * avg_output_std)\n    # print intermediate results\n    print(\n        f\"[Epoch {e:3d}] \"\n        f\"Loss = {avg_loss:.2f} | \"\n        f\"Collapse Level: {collapse_level:.2f} / 1.00\"\n    )"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "To embed the images in the dataset we simply iterate over the test dataloader\nand feed the images to the model backbone. Make sure to disable gradients for\nthis part.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "embeddings = []\nfilenames = []\n\n# disable gradients for faster calculations\nmodel.eval()\nwith torch.no_grad():\n    for i, (x, _, fnames) in enumerate(dataloader_test):\n        # move the images to the gpu\n        x = x.to(device)\n        # embed the images with the pre-trained backbone\n        y = model.backbone(x).flatten(start_dim=1)\n        # store the embeddings and filenames in lists\n        embeddings.append(y)\n        filenames = filenames + list(fnames)\n\n# concatenate the embeddings and convert to numpy\nembeddings = torch.cat(embeddings, dim=0)\nembeddings = embeddings.cpu().numpy()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Scatter Plot and Nearest Neighbors\nNow that we have the embeddings, we can visualize the data with a scatter plot.\nFurther down, we also check out the nearest neighbors of a few example images.\n\nAs a first step, we make a few additional imports.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# for plotting\nimport os\n\nimport matplotlib.offsetbox as osb\nimport matplotlib.pyplot as plt\n\n# for resizing images to thumbnails\nimport torchvision.transforms.functional as functional\nfrom matplotlib import rcParams as rcp\nfrom PIL import Image\n\n# for clustering and 2d representations\nfrom sklearn import random_projection"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Then, we transform the embeddings using UMAP and rescale them to fit in the\n[0, 1] square.\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# for the scatter plot we want to transform the images to a two-dimensional\n# vector space using a random Gaussian projection\nprojection = random_projection.GaussianRandomProjection(n_components=2)\nembeddings_2d = projection.fit_transform(embeddings)\n\n# normalize the embeddings to fit in the [0, 1] square\nM = np.max(embeddings_2d, axis=0)\nm = np.min(embeddings_2d, axis=0)\nembeddings_2d = (embeddings_2d - m) / (M - m)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Let's start with a nice scatter plot of our dataset! The helper function\nbelow will create one.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "def get_scatter_plot_with_thumbnails():\n    \"\"\"Creates a scatter plot with image overlays.\"\"\"\n    # initialize empty figure and add subplot\n    fig = plt.figure()\n    fig.suptitle(\"Scatter Plot of the Sentinel-2 Dataset\")\n    ax = fig.add_subplot(1, 1, 1)\n    # shuffle images and find out which images to show\n    shown_images_idx = []\n    shown_images = np.array([[1.0, 1.0]])\n    iterator = [i for i in range(embeddings_2d.shape[0])]\n    np.random.shuffle(iterator)\n    for i in iterator:\n        # only show image if it is sufficiently far away from the others\n        dist = np.sum((embeddings_2d[i] - shown_images) ** 2, 1)\n        if np.min(dist) < 2e-3:\n            continue\n        shown_images = np.r_[shown_images, [embeddings_2d[i]]]\n        shown_images_idx.append(i)\n\n    # plot image overlays\n    for idx in shown_images_idx:\n        thumbnail_size = int(rcp[\"figure.figsize\"][0] * 2.0)\n        path = os.path.join(path_to_data, filenames[idx])\n        img = Image.open(path)\n        img = functional.resize(img, thumbnail_size)\n        img = np.array(img)\n        img_box = osb.AnnotationBbox(\n            osb.OffsetImage(img, cmap=plt.cm.gray_r),\n            embeddings_2d[idx],\n            pad=0.2,\n        )\n        ax.add_artist(img_box)\n\n    # set aspect ratio\n    ratio = 1.0 / ax.get_data_ratio()\n    ax.set_aspect(ratio, adjustable=\"box\")\n\n\n# get a scatter plot with thumbnail overlays\nget_scatter_plot_with_thumbnails()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Next, we plot example images and their nearest neighbors (calculated from the\nembeddings generated above). This is a very simple approach to find more images\nof a certain type where a few examples are already available. For example,\nwhen a subset of the data is already labelled and one class of images is clearly\nunderrepresented, one can easily query more images of this class from the\nunlabelled dataset.\n\nLet's get to work! The plots are shown below.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "example_images = [\n    \"S2B_MSIL1C_20200526T101559_N0209_R065_T31TGE/tile_00154.png\",  # water 1\n    \"S2B_MSIL1C_20200526T101559_N0209_R065_T32SLJ/tile_00527.png\",  # water 2\n    \"S2B_MSIL1C_20200526T101559_N0209_R065_T32TNL/tile_00556.png\",  # land\n    \"S2B_MSIL1C_20200526T101559_N0209_R065_T31SGD/tile_01731.png\",  # clouds 1\n    \"S2B_MSIL1C_20200526T101559_N0209_R065_T32SMG/tile_00238.png\",  # clouds 2\n]\n\n\ndef get_image_as_np_array(filename: str):\n    \"\"\"Loads the image with filename and returns it as a numpy array.\"\"\"\n    img = Image.open(filename)\n    return np.asarray(img)\n\n\ndef get_image_as_np_array_with_frame(filename: str, w: int = 5):\n    \"\"\"Returns an image as a numpy array with a black frame of width w.\"\"\"\n    img = get_image_as_np_array(filename)\n    ny, nx, _ = img.shape\n    # create an empty image with padding for the frame\n    framed_img = np.zeros((w + ny + w, w + nx + w, 3))\n    framed_img = framed_img.astype(np.uint8)\n    # put the original image in the middle of the new one\n    framed_img[w:-w, w:-w] = img\n    return framed_img\n\n\ndef plot_nearest_neighbors_3x3(example_image: str, i: int):\n    \"\"\"Plots the example image and its eight nearest neighbors.\"\"\"\n    n_subplots = 9\n    # initialize empty figure\n    fig = plt.figure()\n    fig.suptitle(f\"Nearest Neighbor Plot {i + 1}\")\n    #\n    example_idx = filenames.index(example_image)\n    # get distances to the cluster center\n    distances = embeddings - embeddings[example_idx]\n    distances = np.power(distances, 2).sum(-1).squeeze()\n    # sort indices by distance to the center\n    nearest_neighbors = np.argsort(distances)[:n_subplots]\n    # show images\n    for plot_offset, plot_idx in enumerate(nearest_neighbors):\n        ax = fig.add_subplot(3, 3, plot_offset + 1)\n        # get the corresponding filename\n        fname = os.path.join(path_to_data, filenames[plot_idx])\n        if plot_offset == 0:\n            ax.set_title(f\"Example Image\")\n            plt.imshow(get_image_as_np_array_with_frame(fname))\n        else:\n            plt.imshow(get_image_as_np_array(fname))\n        # let's disable the axis\n        plt.axis(\"off\")\n\n\n# show example images for each cluster\nfor i, example_image in enumerate(example_images):\n    plot_nearest_neighbors_3x3(example_image, i)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Next Steps\n\nInterested in exploring other self-supervised models? Check out our other\ntutorials:\n\n- `lightly-moco-tutorial-2`\n- `lightly-simclr-tutorial-3`\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.12.9"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}