self_attention_cv.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "self-attention-cv.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "authorship_tag": "ABX9TyMHPhWUipt1dnFcOkC8OGfh",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/RajkumarGalaxy/ML-Image-Processing/blob/master/self_attention_cv.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DXybxfSTjhYS"
      },
      "source": [
        "# Self-Attention Computer Vision\r\n",
        "\r\n",
        "This notebook has reference from the following sources and papers\r\n",
        "\r\n",
        "https://github.com/The-AI-Summer/self-attention-cv\r\n",
        "https://arxiv.org/pdf/1706.03762.pdf\r\n",
        "https://analyticsindiamag.com/going-beyond-cnn-stand-alone-self-attention/\r\n",
        "https://arxiv.org/pdf/2101.11605"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "1FWhagwLjdzo",
        "outputId": "827cfb50-d31f-4cb0-dca6-ef146feeb369"
      },
      "source": [
        "!pip install self_attention_cv"
      ],
      "execution_count": 3,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Requirement already satisfied: self_attention_cv in /usr/local/lib/python3.7/dist-packages (1.1.0)\n",
            "Requirement already satisfied: torchvision>=0.8 in /usr/local/lib/python3.7/dist-packages (from self_attention_cv) (0.9.0+cu101)\n",
            "Requirement already satisfied: einops>=0.3 in /usr/local/lib/python3.7/dist-packages (from self_attention_cv) (0.3.0)\n",
            "Requirement already satisfied: torch>=1.7 in /usr/local/lib/python3.7/dist-packages (from self_attention_cv) (1.8.0+cu101)\n",
            "Requirement already satisfied: numpy>=1.19 in /usr/local/lib/python3.7/dist-packages (from self_attention_cv) (1.19.5)\n",
            "Requirement already satisfied: pytest>=6.2 in /usr/local/lib/python3.7/dist-packages (from self_attention_cv) (6.2.2)\n",
            "Requirement already satisfied: pillow>=4.1.1 in /usr/local/lib/python3.7/dist-packages (from torchvision>=0.8->self_attention_cv) (7.0.0)\n",
            "Requirement already satisfied: typing-extensions in /usr/local/lib/python3.7/dist-packages (from torch>=1.7->self_attention_cv) (3.7.4.3)\n",
            "Requirement already satisfied: py>=1.8.2 in /usr/local/lib/python3.7/dist-packages (from pytest>=6.2->self_attention_cv) (1.10.0)\n",
            "Requirement already satisfied: importlib-metadata>=0.12; python_version < \"3.8\" in /usr/local/lib/python3.7/dist-packages (from pytest>=6.2->self_attention_cv) (3.7.0)\n",
            "Requirement already satisfied: pluggy<1.0.0a1,>=0.12 in /usr/local/lib/python3.7/dist-packages (from pytest>=6.2->self_attention_cv) (0.13.1)\n",
            "Requirement already satisfied: packaging in /usr/local/lib/python3.7/dist-packages (from pytest>=6.2->self_attention_cv) (20.9)\n",
            "Requirement already satisfied: attrs>=19.2.0 in /usr/local/lib/python3.7/dist-packages (from pytest>=6.2->self_attention_cv) (20.3.0)\n",
            "Requirement already satisfied: toml in /usr/local/lib/python3.7/dist-packages (from pytest>=6.2->self_attention_cv) (0.10.2)\n",
            "Requirement already satisfied: iniconfig in /usr/local/lib/python3.7/dist-packages (from pytest>=6.2->self_attention_cv) (1.1.1)\n",
            "Requirement already satisfied: zipp>=0.5 in /usr/local/lib/python3.7/dist-packages (from importlib-metadata>=0.12; python_version < \"3.8\"->pytest>=6.2->self_attention_cv) (3.4.1)\n",
            "Requirement already satisfied: pyparsing>=2.0.2 in /usr/local/lib/python3.7/dist-packages (from packaging->pytest>=6.2->self_attention_cv) (2.4.7)\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mXNuftZNtnZY"
      },
      "source": [
        "## Multi-head Self Attention"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "aBIJkMtzkDN7",
        "outputId": "23ff19f2-be84-4d22-fc29-73e0a3b1d516"
      },
      "source": [
        "import torch\r\n",
        "from self_attention_cv import MultiHeadSelfAttention\r\n",
        "\r\n",
        "model = MultiHeadSelfAttention(dim=64)\r\n",
        "x = torch.rand(16, 10, 64)  # [batch, tokens, dim]\r\n",
        "mask = torch.zeros(10, 10)  # tokens X tokens\r\n",
        "mask[5:8, 5:8] = 1\r\n",
        "y = model(x, mask)\r\n",
        "\r\n",
        "print('Shape of output is: ', y.shape)\r\n",
        "print('-'*70)\r\n",
        "print('Output corresponding to the first token/patch in the first batch \\n')\r\n",
        "print(y.detach().numpy()[0][0])\r\n"
      ],
      "execution_count": 2,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Shape of output is:  torch.Size([16, 10, 64])\n",
            "----------------------------------------------------------------------\n",
            "Output corresponding to the first token/patch in the first batch \n",
            "\n",
            "[ 0.05095745 -0.3574398   0.2571603   0.01686804 -0.32707906  0.16551468\n",
            " -0.03029712  0.17124459 -0.11935965 -0.05426797 -0.27412418 -0.5849673\n",
            "  0.03273072 -0.07591394  0.0918312   0.14555283 -0.01222516  0.12830907\n",
            " -0.12161413  0.22703463  0.31742722 -0.15440758  0.16580498 -0.1876549\n",
            " -0.14389239 -0.31196266  0.10959084  0.32506582  0.34819284 -0.20865689\n",
            " -0.15298392 -0.02287815 -0.28265458  0.39987215  0.09144013  0.5114911\n",
            " -0.10702722 -0.17274295 -0.03052047 -0.32069373 -0.24848755 -0.19961475\n",
            " -0.25558165  0.16457602  0.04435449  0.14818098 -0.25300896  0.19681376\n",
            " -0.14411886  0.0622221  -0.2663249  -0.08401316 -0.266905   -0.14300469\n",
            "  0.02286307 -0.13363452 -0.00669263 -0.08034971 -0.09160165  0.19021185\n",
            " -0.10703149 -0.09570458  0.07904153 -0.09355289]\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "s2Mw4OMatsT-"
      },
      "source": [
        "## Axial Attention"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "0TljeBmWkGnP",
        "outputId": "ec5d0912-5f24-481c-d361-c0f1a8f3a7eb"
      },
      "source": [
        "# Axial Attention\r\n",
        "from self_attention_cv import AxialAttentionBlock\r\n",
        "\r\n",
        "model = AxialAttentionBlock(in_channels=256, dim=64, heads=8)\r\n",
        "x = torch.rand(1, 256, 64, 64)  # [batch, tokens, dim, dim]\r\n",
        "y = model(x)\r\n",
        "\r\n",
        "print('Shape of output is: ', y.shape)\r\n",
        "print('-'*70)\r\n",
        "print('Output corresponding to the first token/patch in the first batch \\n')\r\n",
        "print(y.detach().numpy()[0][0])"
      ],
      "execution_count": 3,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Shape of output is:  torch.Size([1, 256, 64, 64])\n",
            "----------------------------------------------------------------------\n",
            "Output corresponding to the first token/patch in the first batch \n",
            "\n",
            "[[0.22185817 1.8086953  1.3097283  ... 0.630754   1.1579583  0.        ]\n",
            " [0.         0.         0.         ... 3.1923792  1.9478554  3.044509  ]\n",
            " [1.4585196  1.0503312  0.         ... 0.18355879 0.05406958 0.        ]\n",
            " ...\n",
            " [0.33815414 0.09240472 4.8054843  ... 0.         0.7096393  1.2665645 ]\n",
            " [0.09514064 0.09698692 1.9458563  ... 1.8430982  0.         1.1634666 ]\n",
            " [0.         1.0365889  1.2793136  ... 0.6040628  0.10727388 1.4690821 ]]\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JtOkveV2txCA"
      },
      "source": [
        "## Bottleneck Attention"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "OLXdM48Bs7sX",
        "outputId": "e97a34d0-c0a7-4bd8-9c66-6a8b77751771"
      },
      "source": [
        "from self_attention_cv.bottleneck_transformer import BottleneckBlock\r\n",
        "x = torch.rand(1, 512, 32, 32)\r\n",
        "bottleneck_block = BottleneckBlock(in_channels=512, fmap_size=(32, 32), heads=4, out_channels=1024, pooling=True)\r\n",
        "y = bottleneck_block(x)\r\n",
        "\r\n",
        "print('Shape of output is: ', y.shape)\r\n",
        "print('-'*70)\r\n",
        "print('Output corresponding to the first patch in the first head, first batch \\n')\r\n",
        "print(y.detach().numpy()[0][0][0])"
      ],
      "execution_count": 13,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "forward.. torch.Size([1, 512, 32, 32])\n",
            "Shape of output is:  torch.Size([1, 1024, 16, 16])\n",
            "----------------------------------------------------------------------\n",
            "Output corresponding to the first patch in the first head, first batch \n",
            "\n",
            "[ 2.483006   -0.2608691   0.8990649   1.7123606   0.3083073   0.57465255\n",
            " -1.069338    2.2092774   0.37177175  1.4038159   1.6982754   0.44549412\n",
            " -0.93696225 -0.86148393  1.0227227  -0.33026367]\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QvBX6luZt6mF"
      },
      "source": [
        "## Transformer Encoder\r\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "oRWIv_yqryHO",
        "outputId": "b871e426-4f73-401a-e5e7-90244725d3fc"
      },
      "source": [
        "# Transformer Encoder\r\n",
        "from self_attention_cv import TransformerEncoder\r\n",
        "\r\n",
        "model = TransformerEncoder(dim=64,blocks=6,heads=8)\r\n",
        "x = torch.rand(16, 10, 64)  # [batch, tokens, dim]\r\n",
        "mask = torch.zeros(10, 10)  # tokens X tokens\r\n",
        "mask[5:8, 5:8] = 1\r\n",
        "y = model(x,mask)\r\n",
        "\r\n",
        "print('Shape of output is: ', y.shape)\r\n",
        "print('-'*70)\r\n",
        "print('Output corresponding to the first token/patch in the first batch \\n')\r\n",
        "print(y.detach().numpy()[0][0])"
      ],
      "execution_count": 5,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Shape of output is:  torch.Size([16, 10, 64])\n",
            "----------------------------------------------------------------------\n",
            "Output corresponding to the first token/patch in the first batch \n",
            "\n",
            "[ 1.1898813   1.8956766   1.2114029   0.12728807 -0.3628582  -1.1008008\n",
            " -1.3719841   0.39440683  0.51728094  1.4535215  -3.0603547   0.3478451\n",
            " -0.03671097  1.7242503   0.6375115  -0.11810965 -0.05245292 -0.5992121\n",
            " -0.53516376 -1.2852938  -0.5251507  -0.05185176 -1.1167455   0.14817092\n",
            "  0.85586274  0.9737668   0.08120837 -0.75146204 -0.7960583  -0.8653657\n",
            "  1.705025   -1.2222375   1.1493846  -0.23083282  1.9175692   0.51956785\n",
            " -0.7669539  -1.4143958   0.8146387  -0.7686734  -0.02727643 -1.1802992\n",
            "  0.4200065   0.858947    0.9608477  -0.57810295  1.3265743  -1.1821389\n",
            "  0.6577277  -1.7220386   1.061455    1.4776539  -0.5479512  -1.0370749\n",
            "  0.03364755 -0.25013465  0.34635526 -0.32273477 -0.3752242  -0.5290259\n",
            "  0.71639156  0.5257713  -0.03652823 -1.2284385 ]\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Guu4CBXat-5_"
      },
      "source": [
        "## Vision Transformer "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "MwV5vwsDsRR-",
        "outputId": "fee19029-5e0c-4f23-cc8a-009122ddd380"
      },
      "source": [
        "from self_attention_cv import ViT\r\n",
        "\r\n",
        "model = ViT(img_dim=256, in_channels=3, patch_dim=16, num_classes=10,dim=512)\r\n",
        "x = torch.rand(2, 3, 256, 256)\r\n",
        "y = model(x) # [2,10]\r\n",
        "\r\n",
        "print('Shape of output is: ', y.shape)\r\n",
        "print('-'*70)\r\n",
        "print('Output corresponding to the first image \\n')\r\n",
        "print(y.detach().numpy()[0])"
      ],
      "execution_count": 5,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Shape of output is:  torch.Size([2, 10])\n",
            "----------------------------------------------------------------------\n",
            "Output corresponding to the first image \n",
            "\n",
            "[-0.2508071  -0.0645473  -0.13877457 -0.18308383 -0.3567136   0.34362894\n",
            " -0.35030687  0.3489043   0.5285294   0.08138363]\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZX0DmeXVuH3f"
      },
      "source": [
        "## Vision Transformer with ResNet50"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "B98nDL3bsU81",
        "outputId": "19e108b8-0b6f-4567-f307-55952d30d722"
      },
      "source": [
        "from self_attention_cv import ResNet50ViT\r\n",
        "\r\n",
        "model = ResNet50ViT(img_dim=256, pretrained_resnet=False, \r\n",
        "                        blocks=6, num_classes=10, \r\n",
        "                        dim_linear_block=256, dim=256)\r\n",
        "x = torch.rand(2, 3, 256, 256)\r\n",
        "y = model(x) # [2,10]\r\n",
        "\r\n",
        "print('Shape of output is: ', y.shape)\r\n",
        "print('-'*70)\r\n",
        "print('Output corresponding to the first image \\n')\r\n",
        "print(y.detach().numpy()[0])"
      ],
      "execution_count": 6,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Shape of output is:  torch.Size([2, 10])\n",
            "----------------------------------------------------------------------\n",
            "Output corresponding to the first image \n",
            "\n",
            "[ 0.18863088  0.60527444  0.2697486   0.12503052  0.44111896 -0.2981342\n",
            "  0.5278659  -0.769126   -0.8882933  -1.085641  ]\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JeRk_Nw1uMUY"
      },
      "source": [
        "## TransUNet "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "46U88N95ssVz",
        "outputId": "29bb2106-e3cf-4540-f2e2-2e08e2eff40e"
      },
      "source": [
        "from self_attention_cv.transunet import TransUnet\r\n",
        "x = torch.rand(2, 3, 128, 128)\r\n",
        "model = TransUnet(in_channels=3, img_dim=128, vit_blocks=8,\r\n",
        "vit_dim_linear_mhsa_block=512, classes=5)\r\n",
        "y = model(x) # [2, 5, 128, 128]\r\n",
        "\r\n",
        "print('Shape of output is: ', y.shape)\r\n",
        "print('-'*70)\r\n",
        "print('Output corresponding to the first image \\n')\r\n",
        "print(y.detach().numpy()[0][0])"
      ],
      "execution_count": 7,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Shape of output is:  torch.Size([2, 5, 128, 128])\n",
            "----------------------------------------------------------------------\n",
            "Output corresponding to the first image \n",
            "\n",
            "[[-0.5679928  -0.34442404 -0.27711168 ... -0.5603674  -0.7525164\n",
            "  -0.47777802]\n",
            " [ 0.26423207 -0.00429824 -0.2819016  ... -0.19415903 -0.06176922\n",
            "  -0.64677024]\n",
            " [-0.21234514 -0.3868027  -0.11079104 ... -0.44116968 -0.5203699\n",
            "  -0.5179628 ]\n",
            " ...\n",
            " [ 0.3248333  -0.96315455 -1.3502424  ... -0.3086787  -0.92315567\n",
            "  -0.42170545]\n",
            " [-0.04411215 -0.9422567  -1.0312608  ... -0.8501636  -0.92981017\n",
            "  -0.44236282]\n",
            " [ 0.00445503 -0.27653238 -0.32638735 ... -0.34467506 -0.33213058\n",
            "  -0.67374027]]\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cdIe6dDJudyK"
      },
      "source": [
        "## 1D Absolute Positional Embedding"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "HrY4aR90s9Bl",
        "outputId": "9ad28a4f-bb12-4118-f180-41d8e18dc1e1"
      },
      "source": [
        "from self_attention_cv.pos_embeddings import AbsPosEmb1D\r\n",
        "\r\n",
        "model = AbsPosEmb1D(tokens=20, dim_head=64)\r\n",
        "# batch heads tokens dim_head\r\n",
        "x = torch.rand(2, 3, 20, 64)\r\n",
        "y = model(x)\r\n",
        "\r\n",
        "print('Shape of output is: ', y.shape)\r\n",
        "print('-'*70)\r\n",
        "print('Output corresponding to the first token in the first head, first batch \\n')\r\n",
        "print(y.detach().numpy()[0][0][0])"
      ],
      "execution_count": 8,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Shape of output is:  torch.Size([2, 3, 20, 20])\n",
            "----------------------------------------------------------------------\n",
            "Output corresponding to the first token in the first head, first batch \n",
            "\n",
            "[-0.11605695  0.6117866  -0.1672374   0.43876237  0.23604953  0.1562004\n",
            "  0.48934904  0.12783262  0.70804363  0.50528526 -0.11475162  1.131989\n",
            "  0.38430238 -0.4284703   0.7435987  -0.5122682  -0.36132246 -0.91800463\n",
            "  1.5392176  -0.05165316]\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lE0cNWoruasQ"
      },
      "source": [
        "## 1D Relative Positional Embedding"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ZvBJDQbbtP34",
        "outputId": "2dd258b2-2889-4208-d0da-e073998e352a"
      },
      "source": [
        "from self_attention_cv.pos_embeddings import RelPosEmb1D\r\n",
        "\r\n",
        "model = RelPosEmb1D(tokens=20, dim_head=64, heads=3)\r\n",
        "x = torch.rand(2, 3, 20, 64)\r\n",
        "y = model(x)\r\n",
        "\r\n",
        "print('Shape of output is: ', y.shape)\r\n",
        "print('-'*70)\r\n",
        "print('Output corresponding to the first token in the first head, first batch \\n')\r\n",
        "print(y.detach().numpy()[0][0][0])"
      ],
      "execution_count": 9,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Shape of output is:  torch.Size([2, 3, 20, 20])\n",
            "----------------------------------------------------------------------\n",
            "Output corresponding to the first token in the first head, first batch \n",
            "\n",
            "[-0.5457531   0.0245374   1.206948    1.1106241   0.1697544  -0.04933987\n",
            "  0.24826212  0.38298082 -0.41222885 -0.19180176  0.8518995  -0.63441426\n",
            " -0.9155214   1.2972814   0.52583534  0.05831212 -0.14253221 -0.10725353\n",
            " -0.68426543  0.09718338]\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "EkaDlPfPusR7"
      },
      "source": [
        "## 2D Relative Positional Embedding"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "1uw5PnFbtV1T",
        "outputId": "17acc400-cf35-44f4-9675-44a9f579bdce"
      },
      "source": [
        "from self_attention_cv.pos_embeddings import RelPosEmb2D\r\n",
        "dim = 32  # spatial dim of the feat map\r\n",
        "model = RelPosEmb2D(\r\n",
        "    feat_map_size=(dim, dim),\r\n",
        "    dim_head=128)\r\n",
        "\r\n",
        "x = torch.rand(2, 4, dim*dim, 128)\r\n",
        "y = model(x)\r\n",
        "\r\n",
        "print('Shape of output is: ', y.shape)\r\n",
        "print('-'*70)\r\n",
        "print('Output corresponding to the first patch in the first head, first batch \\n')\r\n",
        "print(y.detach().numpy()[0][0][0])"
      ],
      "execution_count": 11,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Shape of output is:  torch.Size([2, 4, 1024, 1024])\n",
            "----------------------------------------------------------------------\n",
            "Output corresponding to the first patch in the first head, first batch \n",
            "\n",
            "[ 1.7701193  -0.3390236  -0.06327437 ...  0.24196783  0.95375854\n",
            "  0.74266154]\n"
          ],
          "name": "stdout"
        }
      ]
    }
  ]
}