{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "\n# Tutorial 3: Train SimCLR on Clothing\n\nIn this tutorial, we will train a SimCLR model using lightly. The model,\naugmentations and training procedure is from\n[A Simple Framework for Contrastive Learning of Visual Representations](https://arxiv.org/abs/2002.05709).\n\nThe paper explores a rather simple training procedure for contrastive learning.\nSince we use the typical contrastive learning loss based on NCE the method\ngreatly benefits from having larger batch sizes. In this example, we use a batch\nsize of 256 and paired with the input resolution per image of 64x64 pixels and\na resnet-18 model this example requires 16GB of GPU memory.\n\nWe use the\n[clothing dataset from Alex Grigorev](https://github.com/alexeygrigorev/clothing-dataset)\nfor this tutorial.\n\nIn this tutorial you will learn:\n\n- How to create a SimCLR model\n\n- How to generate image representations\n\n- How different augmentations impact the learned representations\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 os\n\nimport matplotlib.pyplot as plt\nimport numpy as np\nimport pytorch_lightning as pl\nimport torch\nimport torch.nn as nn\nimport torchvision\nfrom PIL import Image\nfrom sklearn.neighbors import NearestNeighbors\nfrom sklearn.preprocessing import normalize\n\nfrom lightly.data import LightlyDataset\nfrom lightly.transforms import SimCLRTransform, utils" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Configuration\n\nWe set some configuration parameters for our experiment.\nFeel free to change them and analyze the effect.\n\nThe default configuration with a batch size of 256 and input resolution of 128\nrequires 6GB of GPU memory.\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "num_workers = 8\nbatch_size = 256\nseed = 1\nmax_epochs = 20\ninput_size = 128\nnum_ftrs = 32" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Let's set the seed for our experiments\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "pl.seed_everything(seed)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Make sure `path_to_data` points to the downloaded clothing dataset.\nYou can download it using\n`git clone https://github.com/alexeygrigorev/clothing-dataset.git`\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "path_to_data = \"/datasets/clothing-dataset/images\"" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Setup data augmentations and loaders\n\nThe images from the dataset have been taken from above when the clothing was\non a table, bed or floor. Therefore, we can make use of additional augmentations\nsuch as vertical flip or random rotation (90 degrees).\nBy adding these augmentations we learn our model invariance regarding the\norientation of the clothing piece. E.g. we don't care if a shirt is upside down\nbut more about the structure which make it a shirt.\n\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, vf_prob=0.5, rr_prob=0.5)\n\n# We create a torchvision transformation for embedding the dataset after\n# training\ntest_transform = 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\ndataset_train_simclr = LightlyDataset(input_dir=path_to_data, transform=transform)\n\ndataset_test = LightlyDataset(input_dir=path_to_data, transform=test_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)\n\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 SimCLR Model\nNow we create the SimCLR model. We implement it as a PyTorch Lightning Module\nand use a ResNet-18 backbone from Torchvision. Lightly provides implementations\nof the SimCLR projection head and loss function in the `SimCLRProjectionHead`\nand `NTXentLoss` classes. We can simply import them and combine the building\nblocks in the module.\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "from lightly.loss import NTXentLoss\nfrom lightly.models.modules.heads import SimCLRProjectionHead\n\n\nclass SimCLRModel(pl.LightningModule):\n def __init__(self):\n super().__init__()\n\n # create a ResNet backbone and remove the classification head\n resnet = torchvision.models.resnet18()\n self.backbone = nn.Sequential(*list(resnet.children())[:-1])\n\n hidden_dim = resnet.fc.in_features\n self.projection_head = SimCLRProjectionHead(hidden_dim, hidden_dim, 128)\n\n self.criterion = NTXentLoss()\n\n def forward(self, x):\n h = self.backbone(x).flatten(start_dim=1)\n z = self.projection_head(h)\n return z\n\n def training_step(self, batch, batch_idx):\n (x0, x1), _, _ = batch\n z0 = self.forward(x0)\n z1 = self.forward(x1)\n loss = self.criterion(z0, z1)\n self.log(\"train_loss_ssl\", loss)\n return loss\n\n def configure_optimizers(self):\n optim = torch.optim.SGD(\n self.parameters(), lr=6e-2, momentum=0.9, weight_decay=5e-4\n )\n scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optim, max_epochs)\n return [optim], [scheduler]" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Train the module using the PyTorch Lightning Trainer on a single GPU.\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "model = SimCLRModel()\ntrainer = pl.Trainer(max_epochs=max_epochs, devices=1, accelerator=\"gpu\")\ntrainer.fit(model, dataloader_train_simclr)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Next we create a helper function to generate embeddings\nfrom our test images using the model we just trained.\nNote that only the backbone is needed to generate embeddings,\nthe projection head is only required for the training.\nMake sure to put the model into eval mode for this part!\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "def generate_embeddings(model, dataloader):\n \"\"\"Generates representations for all images in the dataloader with\n the given model\n \"\"\"\n embeddings = []\n filenames = []\n with torch.no_grad():\n for img, _, fnames in dataloader:\n img = img.to(model.device)\n emb = model.backbone(img).flatten(start_dim=1)\n embeddings.append(emb)\n filenames.extend(fnames)\n\n embeddings = torch.cat(embeddings, 0)\n embeddings = normalize(embeddings)\n return embeddings, filenames\n\n\nmodel.eval()\nembeddings, filenames = generate_embeddings(model, dataloader_test)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Visualize Nearest Neighbors\nLet's look at the trained embedding and visualize the nearest neighbors for\na few random samples.\n\nWe create some helper functions to simplify the work\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "def get_image_as_np_array(filename: str):\n \"\"\"Returns an image as an numpy array\"\"\"\n img = Image.open(filename)\n return np.asarray(img)\n\n\ndef plot_knn_examples(embeddings, filenames, n_neighbors=3, num_examples=6):\n \"\"\"Plots multiple rows of random images with their nearest neighbors\"\"\"\n # lets look at the nearest neighbors for some samples\n # we use the sklearn library\n nbrs = NearestNeighbors(n_neighbors=n_neighbors).fit(embeddings)\n distances, indices = nbrs.kneighbors(embeddings)\n\n # get 5 random samples\n samples_idx = np.random.choice(len(indices), size=num_examples, replace=False)\n\n # loop through our randomly picked samples\n for idx in samples_idx:\n fig = plt.figure()\n # loop through their nearest neighbors\n for plot_x_offset, neighbor_idx in enumerate(indices[idx]):\n # add the subplot\n ax = fig.add_subplot(1, len(indices[idx]), plot_x_offset + 1)\n # get the corresponding filename for the current index\n fname = os.path.join(path_to_data, filenames[neighbor_idx])\n # plot the image\n plt.imshow(get_image_as_np_array(fname))\n # set the title to the distance of the neighbor\n ax.set_title(f\"d={distances[idx][plot_x_offset]:.3f}\")\n # let's disable the axis\n plt.axis(\"off\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Let's do the plot of the images. The leftmost image is the query image whereas\nthe ones next to it on the same row are the nearest neighbors.\nIn the title we see the distance of the neighbor.\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "plot_knn_examples(embeddings, filenames)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Color Invariance\nLet's train again without color augmentation. This will force our model to\nrespect the colors in the images.\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "# Set color jitter and gray scale probability to 0\nnew_transform = SimCLRTransform(\n input_size=input_size, vf_prob=0.5, rr_prob=0.5, cj_prob=0.0, random_gray_scale=0.0\n)\n\n# let's update the transform on the training dataset\ndataset_train_simclr.transform = new_transform\n\n# then train a new model\nmodel = SimCLRModel()\ntrainer = pl.Trainer(max_epochs=max_epochs, devices=1, accelerator=\"gpu\")\ntrainer.fit(model, dataloader_train_simclr)\n\n# and generate again embeddings from the test set\nmodel.eval()\nembeddings, filenames = generate_embeddings(model, dataloader_test)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "other example\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "plot_knn_examples(embeddings, filenames)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "What's next?\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "# You could use the pre-trained model and train a classifier on top.\npretrained_resnet_backbone = model.backbone\n\n# you can also store the backbone and use it in another code\nstate_dict = {\"resnet18_parameters\": pretrained_resnet_backbone.state_dict()}\ntorch.save(state_dict, \"model.pth\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "THIS COULD BE IN A NEW FILE (e.g. inference.py)\n\nMake sure you place the `model.pth` file in the same folder as this code\n\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "collapsed": false }, "outputs": [], "source": [ "# load the model in a new file for inference\nresnet18_new = torchvision.models.resnet18()\n\n# note that we need to create exactly the same backbone in order to load the weights\nbackbone_new = nn.Sequential(*list(resnet18_new.children())[:-1])\n\nckpt = torch.load(\"model.pth\")\nbackbone_new.load_state_dict(ckpt[\"resnet18_parameters\"])" ] }, { "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-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.12.9" } }, "nbformat": 4, "nbformat_minor": 0 }