{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import torch.nn as nn\n",
    "import cv2 as cv\n",
    "from torchvision import models, transforms"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load the model\n",
    "model = models.resnet50(pretrained=True)\n",
    "#print(model)\n",
    "model_weights = [] # we will save the conv layer weights in this list\n",
    "conv_layers = [] # we will save the 49 conv layers in this list\n",
    "# get all the model children as list\n",
    "model_children = list(model.children())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Total convolutional layers: 49\n"
     ]
    }
   ],
   "source": [
    "counter = 0 \n",
    "# append all the conv layers and their respective weights to the list\n",
    "for i in range(len(model_children)):\n",
    "    if type(model_children[i]) == nn.Conv2d:\n",
    "        counter += 1\n",
    "        model_weights.append(model_children[i].weight)\n",
    "        conv_layers.append(model_children[i])\n",
    "    elif type(model_children[i]) == nn.Sequential:\n",
    "        for j in range(len(model_children[i])):\n",
    "            for child in model_children[i][j].children():\n",
    "                if type(child) == nn.Conv2d:\n",
    "                    counter += 1\n",
    "                    model_weights.append(child.weight)\n",
    "                    conv_layers.append(child)\n",
    "print(f\"Total convolutional layers: {counter}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CONV: Conv2d(3, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False) ====> SHAPE: torch.Size([64, 3, 7, 7])\n",
      "CONV: Conv2d(64, 64, kernel_size=(1, 1), stride=(1, 1), bias=False) ====> SHAPE: torch.Size([64, 64, 1, 1])\n",
      "CONV: Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False) ====> SHAPE: torch.Size([64, 64, 3, 3])\n",
      "CONV: Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False) ====> SHAPE: torch.Size([256, 64, 1, 1])\n",
      "CONV: Conv2d(256, 64, kernel_size=(1, 1), stride=(1, 1), bias=False) ====> SHAPE: torch.Size([64, 256, 1, 1])\n",
      "CONV: Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False) ====> SHAPE: torch.Size([64, 64, 3, 3])\n",
      "CONV: Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False) ====> SHAPE: torch.Size([256, 64, 1, 1])\n",
      "CONV: Conv2d(256, 64, kernel_size=(1, 1), stride=(1, 1), bias=False) ====> SHAPE: torch.Size([64, 256, 1, 1])\n",
      "CONV: Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False) ====> SHAPE: torch.Size([64, 64, 3, 3])\n",
      "CONV: Conv2d(64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False) ====> SHAPE: torch.Size([256, 64, 1, 1])\n",
      "CONV: Conv2d(256, 128, kernel_size=(1, 1), stride=(1, 1), bias=False) ====> SHAPE: torch.Size([128, 256, 1, 1])\n",
      "CONV: Conv2d(128, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False) ====> SHAPE: torch.Size([128, 128, 3, 3])\n",
      "CONV: Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False) ====> SHAPE: torch.Size([512, 128, 1, 1])\n",
      "CONV: Conv2d(512, 128, kernel_size=(1, 1), stride=(1, 1), bias=False) ====> SHAPE: torch.Size([128, 512, 1, 1])\n",
      "CONV: Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False) ====> SHAPE: torch.Size([128, 128, 3, 3])\n",
      "CONV: Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False) ====> SHAPE: torch.Size([512, 128, 1, 1])\n",
      "CONV: Conv2d(512, 128, kernel_size=(1, 1), stride=(1, 1), bias=False) ====> SHAPE: torch.Size([128, 512, 1, 1])\n",
      "CONV: Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False) ====> SHAPE: torch.Size([128, 128, 3, 3])\n",
      "CONV: Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False) ====> SHAPE: torch.Size([512, 128, 1, 1])\n",
      "CONV: Conv2d(512, 128, kernel_size=(1, 1), stride=(1, 1), bias=False) ====> SHAPE: torch.Size([128, 512, 1, 1])\n",
      "CONV: Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False) ====> SHAPE: torch.Size([128, 128, 3, 3])\n",
      "CONV: Conv2d(128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False) ====> SHAPE: torch.Size([512, 128, 1, 1])\n",
      "CONV: Conv2d(512, 256, kernel_size=(1, 1), stride=(1, 1), bias=False) ====> SHAPE: torch.Size([256, 512, 1, 1])\n",
      "CONV: Conv2d(256, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False) ====> SHAPE: torch.Size([256, 256, 3, 3])\n",
      "CONV: Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False) ====> SHAPE: torch.Size([1024, 256, 1, 1])\n",
      "CONV: Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False) ====> SHAPE: torch.Size([256, 1024, 1, 1])\n",
      "CONV: Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False) ====> SHAPE: torch.Size([256, 256, 3, 3])\n",
      "CONV: Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False) ====> SHAPE: torch.Size([1024, 256, 1, 1])\n",
      "CONV: Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False) ====> SHAPE: torch.Size([256, 1024, 1, 1])\n",
      "CONV: Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False) ====> SHAPE: torch.Size([256, 256, 3, 3])\n",
      "CONV: Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False) ====> SHAPE: torch.Size([1024, 256, 1, 1])\n",
      "CONV: Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False) ====> SHAPE: torch.Size([256, 1024, 1, 1])\n",
      "CONV: Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False) ====> SHAPE: torch.Size([256, 256, 3, 3])\n",
      "CONV: Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False) ====> SHAPE: torch.Size([1024, 256, 1, 1])\n",
      "CONV: Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False) ====> SHAPE: torch.Size([256, 1024, 1, 1])\n",
      "CONV: Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False) ====> SHAPE: torch.Size([256, 256, 3, 3])\n",
      "CONV: Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False) ====> SHAPE: torch.Size([1024, 256, 1, 1])\n",
      "CONV: Conv2d(1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False) ====> SHAPE: torch.Size([256, 1024, 1, 1])\n",
      "CONV: Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False) ====> SHAPE: torch.Size([256, 256, 3, 3])\n",
      "CONV: Conv2d(256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False) ====> SHAPE: torch.Size([1024, 256, 1, 1])\n",
      "CONV: Conv2d(1024, 512, kernel_size=(1, 1), stride=(1, 1), bias=False) ====> SHAPE: torch.Size([512, 1024, 1, 1])\n",
      "CONV: Conv2d(512, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False) ====> SHAPE: torch.Size([512, 512, 3, 3])\n",
      "CONV: Conv2d(512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False) ====> SHAPE: torch.Size([2048, 512, 1, 1])\n",
      "CONV: Conv2d(2048, 512, kernel_size=(1, 1), stride=(1, 1), bias=False) ====> SHAPE: torch.Size([512, 2048, 1, 1])\n",
      "CONV: Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False) ====> SHAPE: torch.Size([512, 512, 3, 3])\n",
      "CONV: Conv2d(512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False) ====> SHAPE: torch.Size([2048, 512, 1, 1])\n",
      "CONV: Conv2d(2048, 512, kernel_size=(1, 1), stride=(1, 1), bias=False) ====> SHAPE: torch.Size([512, 2048, 1, 1])\n",
      "CONV: Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False) ====> SHAPE: torch.Size([512, 512, 3, 3])\n",
      "CONV: Conv2d(512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False) ====> SHAPE: torch.Size([2048, 512, 1, 1])\n"
     ]
    }
   ],
   "source": [
    "for weight, conv in zip(model_weights, conv_layers):\n",
    "    # print(f\"WEIGHT: {weight} \\nSHAPE: {weight.shape}\")\n",
    "    print(f\"CONV: {conv} ====> SHAPE: {weight.shape}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 30368 characters]\n",
      "text/plain": [
       "<Figure size 1440x1224 with 64 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# visualize the first conv layer filters\n",
    "plt.figure(figsize=(20, 17))\n",
    "for i, filter in enumerate(model_weights[0]):\n",
    "    plt.subplot(8, 8, i+1) # (8, 8) because in conv0 we have 7x7 filters and total of 64 (see printed shapes)\n",
    "    plt.imshow(filter[0, :, :].detach(), cmap='gray')\n",
    "    plt.axis('off')\n",
    "    plt.savefig('data/outputs/resnet/filter1.png')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([64, 3, 7, 7])"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model_weights[0].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "49"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(model_weights)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 186744 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([3, 224, 224])\n",
      "torch.Size([1, 3, 224, 224])\n"
     ]
    }
   ],
   "source": [
    "# read and visualize an image\n",
    "img = cv.imread(\"data/images/dog.png\")\n",
    "img = cv.cvtColor(img, cv.COLOR_BGR2RGB)\n",
    "plt.imshow(img)\n",
    "plt.show()\n",
    "# define the transforms\n",
    "transform = transforms.Compose([\n",
    "    transforms.ToPILImage(),\n",
    "    transforms.Resize((224, 224)),\n",
    "    transforms.ToTensor(),\n",
    "])\n",
    "img = np.array(img)\n",
    "# apply the transforms\n",
    "img = transform(img)\n",
    "print(img.size())\n",
    "# unsqueeze to add a batch dimension\n",
    "img = img.unsqueeze(0)\n",
    "print(img.size())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "# pass the image through all the layers\n",
    "results = [conv_layers[0](img)]\n",
    "for i in range(1, len(conv_layers)):\n",
    "    # pass the result from the last layer to the next layer\n",
    "    results.append(conv_layers[i](results[-1]))\n",
    "# make a copy of the `results`\n",
    "outputs = results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "#len(results)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Saving layer 0 feature maps...\n",
      "Saving layer 1 feature maps...\n",
      "Saving layer 2 feature maps...\n",
      "Saving layer 3 feature maps...\n",
      "Saving layer 4 feature maps...\n"
     ]
    }
   ],
   "source": [
    "# visualize 36 features from each layer \n",
    "# (although there are more feature maps in the upper layers)\n",
    "all_layers_num = 5   # change the number of length to control how many layers you want to visualize\n",
    "for num_layer in range(all_layers_num): \n",
    "    plt.figure(figsize=(20, 20))\n",
    "    layer_viz = outputs[num_layer][0, :, :, :]\n",
    "    layer_viz = layer_viz.data\n",
    "    #print(layer_viz.size())\n",
    "    for i, filter in enumerate(layer_viz):\n",
    "        if i == 36: # we will visualize only 8x8 blocks from each layer\n",
    "            break\n",
    "        plt.subplot(6, 6, i + 1)\n",
    "        plt.imshow(filter, cmap='gray')\n",
    "        plt.axis(\"off\")\n",
    "    print(f\"Saving layer {num_layer} feature maps...\")\n",
    "    plt.savefig(f\"data/outputs/resnet/layer_{num_layer}.png\")  # change the path to save the feature maps\n",
    "    #plt.show() # use this line to show the figure in jupyter notebook\n",
    "    plt.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.7.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
