Vgg16_Torch_KernelVisualization.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Vgg16_Torch_KernelVisualization.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/dinoelT/YoutubeCode/blob/main/Vgg16_Torch_KernelVisualization.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hMRWcu4p8_80"
      },
      "source": [
        "# Visualization of the patterns the kernels of the neural network are looking for. \r\n",
        "Youtube channel: [Leonid's Dreams](https://www.youtube.com/channel/UCvOlcEWLNSm_JbMJVDhcSyw)"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "M1QHu5elc9cn"
      },
      "source": [
        "import torch\n",
        "import torchvision.models as models\n",
        "import numpy as np\n",
        "from torch import nn, optim\n",
        "import matplotlib.pyplot as plt\n",
        "import torch.nn.functional as F\n",
        "\n",
        "from IPython.display import clear_output, display"
      ],
      "execution_count": 2,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "SmlrErEVyWG2"
      },
      "source": [
        "#Pick a layer to visualize, from conv_1 ... conv_13\r\n",
        "layerName = \"conv_10\"\r\n",
        "\r\n",
        "''' The image will start from the <initialSize>, and will be multiplied with the \r\n",
        "    <scale_factor> for <nrOfScalings> times. In between the scalings, the image\r\n",
        "    will be optimized to minimize the loss.\r\n",
        "    Example: if imageSize is (100,150) and the nrOfScalings=10 with a\r\n",
        "    scaling_factor = 1.2, the final image size will be \r\n",
        "    (100 * 1.2^10, 150 * 1.2^10) = (619, 928)\r\n",
        "'''\r\n",
        "\r\n",
        "#Pick the initial image size, that will be resized \r\n",
        "imageSize = (100, 100)\r\n",
        "\r\n",
        "#Pick the number of times the image will be scaled\r\n",
        "nrOfScalings = 10\r\n",
        "\r\n",
        "#Choose the scaling factor\r\n",
        "scale_factor = 1.2\r\n",
        "\r\n",
        "#Pick the number of filters you want to visualize\r\n",
        "NrOfFiltersToVisualize = 10"
      ],
      "execution_count": 3,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "GNMPopvCdEu7"
      },
      "source": [
        "def getVgg16Model():\n",
        "  \"\"\" Returns a Vgg16 model with AveragePool instead of Maxpool \n",
        "      and with naming convention: avgpool_x, conv_x, relu_x\n",
        "  \"\"\"\n",
        "  initialModel = models.vgg16(pretrained=True)\n",
        "\n",
        "  modifiedModel = torch.nn.modules.container.Sequential()\n",
        "  #Dictionary with the index of the layers for naming\n",
        "  index_dict = {'avgpool': 1, 'conv':1, 'relu':1}\n",
        "\n",
        "  for layer in initialModel.features:\n",
        "    \n",
        "    if isinstance(layer, torch.nn.MaxPool2d):\n",
        "      name = 'avgpool'\n",
        "      i = index_dict[name]\n",
        "      modifiedModel.add_module(name+'_'+str(i), torch.nn.AvgPool2d((2,2)))\n",
        "      index_dict[name] = index_dict[name] +  1\n",
        "    elif isinstance(layer, torch.nn.Conv2d):\n",
        "      name = 'conv'\n",
        "      i = index_dict[name]\n",
        "      modifiedModel.add_module(name+'_'+str(i),layer)\n",
        "      index_dict[name] = index_dict[name] +  1\n",
        "    elif isinstance(layer, torch.nn.ReLU):\n",
        "      name = 'relu'\n",
        "      i = index_dict[name]\n",
        "      modifiedModel.add_module(name+'_'+str(i),layer)\n",
        "      index_dict[name] = index_dict[name] +  1 \n",
        "    else:\n",
        "      break\n",
        "  return modifiedModel\n",
        "\n",
        "\n",
        "def getRandomImage(imageSize):\n",
        "  \"\"\" Returns a random image with shape (1,3,y,x)\n",
        "      imageSize = tuple([y,x])\n",
        "  \"\"\"\n",
        "  img = np.random.uniform(0, 1,(1,*imageSize, 3))\n",
        "  # Normalizing input for vgg16\n",
        "  mean = [0.485, 0.456, 0.406]\n",
        "  std = [0.229, 0.224, 0.225]\n",
        "  img = (img - mean) / std\n",
        "  img = np.transpose(img, axes = (0,3,1,2)).astype(\"float32\")\n",
        "  generatedImage = torch.tensor(img.copy(), requires_grad=True)\n",
        "  return generatedImage\n",
        " \n",
        "\n",
        "def Loss_lpf(featureMap,indexFilterMaximize, generatedImage, alpha):\n",
        "  \"\"\" Calculates the loss with low pass filter\"\"\"\n",
        "  a = torch.sum(featureMap[:,indexFilterMaximize,:,:])\n",
        "  b = lowPassFilterLoss(generatedImage)\n",
        "  loss = b/(a + 0.1)\n",
        "  return loss\n",
        "\n",
        "def Loss(featureMap,indexFilterMaximize, alpha):\n",
        "  \"\"\" Calculates the loss\"\"\"\n",
        "  a = torch.sum(featureMap[:,indexFilterMaximize,:,:])\n",
        "  loss = alpha/(a + 0.1)\n",
        "  return loss\n",
        "\n",
        "def tensorToImage(tensor):\n",
        "  \"\"\" Transforms image tensor into numpy array with shape (y,x,3)\"\"\"\n",
        "  image = tensor.cpu().detach().numpy()[0]\n",
        "  image = np.transpose(image, axes = (1,2,0))\n",
        "  return image\n",
        "\n",
        "def lowPassFilterLoss(inputImg):\n",
        "  \"\"\" Calculates the low pass filter loss on the horizontal and vertical axes.\n",
        "  inputImg size : (1, 3, y, x)\n",
        "  \"\"\"\n",
        "\n",
        "  verticalFilter = torch.sum( torch.abs(inputImg[:,:, :-1, :] - inputImg[:,:, 1:, :] ))\n",
        "  horizontalFilter = torch.sum( torch.abs(inputImg[:,:,:, :-1] - inputImg[:,:,:, 1:] ))\n",
        "\n",
        "  return verticalFilter + horizontalFilter\n",
        "\n",
        "def getMaxFiltersIndexes(imageSize, nrFilters):\n",
        "  \"\"\" Returns the indexes of the kernels that had higher activations when applied\n",
        "  to a random image(if the activations are too low or zero, the gradient is zero)\n",
        "  imageSize = tuple([y,x])\n",
        "  nrFilters - number of filter indexes to be returned, sorted by the value of \n",
        "  the activations (from high to low)\n",
        "  \"\"\"\n",
        "  inputImage = getRandomImage(imageSize).to(device)\n",
        "\n",
        "  out = model(inputImage)\n",
        "\n",
        "  a = features.cpu().detach().numpy()[0]\n",
        "  s = np.sum(a, axis = (1,2))\n",
        "  s_sorted =  np.sort(s)[-nrFilters:]\n",
        "  min_element = np.floor( np.amin(s_sorted))\n",
        "  return np.where(s > min_element)[0]\n"
      ],
      "execution_count": 4,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "s5j0_irWovcr"
      },
      "source": [
        "#handle - holds a reference to the hook we added, in case we want to remove it\n",
        "handle = None\n",
        "\n",
        "#features - this is where the output of the layer will be stored\n",
        "features = None\n",
        "\n",
        "def featureSaveHook(module, input, output):\n",
        "  global features\n",
        "  features = output\n",
        "\n",
        "def addHook(model, layerName):\n",
        "  global handle\n",
        "\n",
        "  #Clear all hooks\n",
        "  if handle is not None:\n",
        "    handle.remove()\n",
        "\n",
        "  layerIndex = list(dict(model.named_children()).keys()).index(layerName)\n",
        "  handle = model[layerIndex].register_forward_hook(featureSaveHook)\n",
        "\n",
        "\n",
        "#Create model\n",
        "model = getVgg16Model()\n",
        "\n",
        "#Create hook for saving features\n",
        "addHook(model, layerName)\n",
        "\n",
        "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
        "print(\"Used device:\", device)\n",
        "model = model.to(device)\n",
        "\n",
        "layerList = [layer for layer in dict(model.named_children()).keys()]\n",
        "print(\"Layer names:\",layerList)\n",
        "\n",
        "#create a folder to save the images in \n",
        "import os\n",
        "\n",
        "path = os.getcwd()\n",
        "os.mkdir(layerName)  "
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "6VQ3FXSd9PHB"
      },
      "source": [
        "filterIndexes = getMaxFiltersIndexes(imageSize, NrOfFiltersToVisualize)\n",
        "print(\"Filter indexes:\",filterIndexes)\n",
        "\n",
        "for filterIndex in filterIndexes:\n",
        "  inputImage = getRandomImage(imageSize)\n",
        "  #print(inputImage.shape)\n",
        "  inputImage = inputImage.to(device).detach().requires_grad_(True)\n",
        "\n",
        "  for i in range(nrOfScalings):\n",
        "    inputImage = F.interpolate(inputImage, scale_factor = scale_factor)\n",
        "    inputImage = inputImage.to(device).detach().requires_grad_(True)\n",
        "    optimizer = optim.Adam([inputImage], lr=0.1)\n",
        "    print(\"\\n<<<<<<<<Iteration\",i,\">>>>>>>>>\")\n",
        "    print(\"Image shape:\", inputImage.shape,\"\\n\")\n",
        "    for epoch in range(30):\n",
        "      featureMaps = model(inputImage)\n",
        "      loss = Loss_lpf(features, filterIndex, inputImage, 10e7)\n",
        "      #loss = Loss(features, filterIndex, 1000)\n",
        "      print(loss.item())\n",
        "      optimizer.zero_grad()\n",
        "      loss.backward()\n",
        "      optimizer.step()\n",
        "      with torch.no_grad():\n",
        "        inputImage[:] = inputImage.clamp(0, 1)\n",
        "  image = tensorToImage(inputImage)\n",
        "  plt.imsave(layerName+\"/\"+layerName + \"_filter\"+str(filterIndex)+\"_lpf.png\", image)\n",
        "  clear_output()\n"
      ],
      "execution_count": 6,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "EmRhJnJ1z2gK"
      },
      "source": [
        "#Save folder as an archive\r\n",
        "import shutil\r\n",
        "shutil.make_archive(layerName, 'zip', layerName)"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}