CNN_Torch.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "view-in-github"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/ShoaibSajid/Python_CNN/blob/main/CNN_Torch.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cw04b-NnPF_e"
      },
      "outputs": [],
      "source": [
        "%cd ~\n",
        "![ ! -d \"./Python_CNN\" ] && echo \"Github Repo DOES NOT exists.\"\n",
        "![ ! -d \"./Python_CNN\" ] && git clone https://github.com/ShoaibSajid/Python_CNN\n",
        "![ -d \"./Python_CNN\" ] && echo \"Github Repo exist.\"\n",
        "%cd Python_CNN\n",
        "%pip install -r requirements.txt\n",
        "!echo \n",
        "!echo List of files\n",
        "!ls"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "I-725DWZTpo6",
        "outputId": "19e9c008-fb83-4b29-f869-e7d8e7c2b63c"
      },
      "outputs": [],
      "source": [
        "# Importing libraries\n",
        "\n",
        "from matplotlib.pyplot import *\n",
        "import torch.nn as nn\n",
        "from tqdm import tqdm\n",
        "import numpy as np\n",
        "import pickle\n",
        "import torch\n",
        "import os\n",
        "import mnist\n",
        "from torch.cuda import amp\n",
        "\n",
        "os.environ[\"CUDA_VISIBLE_DEVICES\"]=\"\"\n",
        "device = \"cpu\"\n",
        "print(f\"Using {device} device\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "iAoYBSKmWKWE"
      },
      "outputs": [],
      "source": [
        "# Define Model\n",
        "class simple_network(nn.Module):\n",
        "    def __init__(self):\n",
        "        super(simple_network, self).__init__()\n",
        "        self.conv0   = nn.Conv2d    (   1,  8,   kernel_size=(3, 3),     bias=False  , padding=1)   \n",
        "        self.max0    = nn.MaxPool2d (            kernel_size=(2,2)                   )   \n",
        "        self.conv1   = nn.Conv2d    (   8,  16,  kernel_size=(3, 3),     bias=False , padding=1 )   \n",
        "        self.max1    = nn.MaxPool2d (            kernel_size=(2,2)                   )   \n",
        "        self.flat    = nn.Flatten   (                                                )   \n",
        "        self.linear0 = nn.Linear    (            784,  784                           )\n",
        "        self.linear1 = nn.Linear    (            784,  10                            )\n",
        "        self.soft    = nn.Softmax   (            dim=1                               )\n",
        "        \n",
        "        self.debug=False\n",
        "        \n",
        "\n",
        "    def forward(self,x):\n",
        "        x=x.to(torch.device(\"cpu\"))\n",
        "        x_conv0      = self.conv0     (   x              )  \n",
        "        x_max0       = self.max0      (   x_conv0        )\n",
        "        x_conv1      = self.conv1     (   x_max0         )  \n",
        "        x_max1       = self.max1      (   x_conv1        )\n",
        "        x_flat       = self.flat      (   x_max1         )\n",
        "        x_linear0    = self.linear0   (   x_flat         )\n",
        "        x_linear1    = self.linear1   (   x_linear0      )\n",
        "        x_prob       = self.soft      (   x_linear1      )\n",
        "\n",
        "        # TODO: Fix Learning Rate and Gradient Descent Method, check batch size\n",
        "        if self.debug:\n",
        "            im=x[0][0]\n",
        "            print(f\"Input Image: {im[-4]}\\n\")\n",
        "            \n",
        "            print(f\"conv0 filters: {model.state_dict()['conv0.weight'][0][0]}\\n\")\n",
        "            \n",
        "            print(f\"x_conv0 : {x_conv0[0,0,:,:][-1]}\\n\")\n",
        "\n",
        "            print(f\"MaxPool0: {x_max0[0,0][-1]}\\n\")\n",
        "            \n",
        "            print(f\"conv1 filters: {model.state_dict()['conv1.weight'][0][0]}\\n\")\n",
        "            \n",
        "            print(f\"x_conv1 : {x_conv1[0,0,:,:][-1]}\\n\")\n",
        "\n",
        "            print(f\"MaxPool1: {x_max1[0,0][-1]}\\n\")\n",
        "            \n",
        "            print(f\"FC0 Weight: {model.state_dict()['linear0.weight'][0][:10]}\\n\")\n",
        "            \n",
        "            print(f\"FC0 Output: {x_linear0[0,:10]}\\n\")\n",
        "            \n",
        "            print(f\"FC1 Weight: {model.state_dict()['linear1.weight'][0][:10]}\\n\")\n",
        "            \n",
        "            print(f\"FC1 Output: {x_linear1[0,:10]}\\n\")\n",
        "            \n",
        "            print(f\"SoftMax Output: {x_prob}\\n\")\n",
        "            \n",
        "        return x_prob"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "mqYZeDDpU8B4",
        "outputId": "6fbe1782-9cb5-4e15-dbf9-b241f5749f28"
      },
      "outputs": [],
      "source": [
        "# Initialize model and print\n",
        "\n",
        "model = simple_network()\n",
        "model.double()\n",
        "print(model)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_YL3A_cuU8B5"
      },
      "outputs": [],
      "source": [
        "# Hyper Parameters\n",
        "\n",
        "n_epochs = 3\n",
        "# batch_size_train = 64\n",
        "# batch_size_test = 1000\n",
        "learning_rate = 0.005\n",
        "momentum = 0\n",
        "log_interval = 1\n",
        "torch.backends.cudnn.enabled = False"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "YMC2OtZsU8B6"
      },
      "outputs": [],
      "source": [
        "# Optimizer\n",
        "optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate,momentum=momentum)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jBCvlbF3U8B6",
        "outputId": "9d624884-2590-40fe-cc07-0a072928acda"
      },
      "outputs": [],
      "source": [
        "# Scaler\n",
        "scaler = amp.GradScaler()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cfnxm2MUU8B9"
      },
      "outputs": [],
      "source": [
        "# Loss Function\n",
        "calculate_loss = nn.CrossEntropyLoss()\n",
        "# calculate_loss_nlll = nn.NLLLoss()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8qQYj2N3U8B_"
      },
      "outputs": [],
      "source": [
        "# Define Dataset\n",
        "train_images = mnist.train_images()\n",
        "train_labels = mnist.train_labels()\n",
        "test_images = mnist.test_images()\n",
        "test_labels = mnist.test_labels()\n",
        "train_images = (train_images / 255) - 0.5\n",
        "test_images = (test_images / 255) - 0.5\n",
        "train_images = np.expand_dims(train_images, axis=3)\n",
        "test_images = np.expand_dims(test_images, axis=3)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "eHQfoTeNU8CC"
      },
      "outputs": [],
      "source": [
        "# Counters to keep track of losses\n",
        "train_losses = []\n",
        "train_counter = []\n",
        "test_losses = []\n",
        "test_counter = [i*len(train_labels) for i in range(n_epochs + 1)]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "SuSyIPjHU8CE"
      },
      "outputs": [],
      "source": [
        "# Function for copying weights from basic model \n",
        "def copy_weights():\n",
        "    # TODO: To copy filter weights\n",
        "    name = \"weights/debug.pkl\"\n",
        "    print(f\"Loading weights from {name}\")\n",
        "    weight_file = open(name, \"rb\")\n",
        "    weights = pickle.load(weight_file)\n",
        "    print(f\"\\nLoading weights from {name} file\")\n",
        "    \n",
        "    fconv0   = weights[\"conv0\"]  \n",
        "    fconv1   = weights[\"conv1\"]  \n",
        "    flinear0 = weights[\"fc0_weights\"]\n",
        "    blinear0 = weights[\"fc0_biases\" ]\n",
        "    flinear1 = weights[\"fc1_weights\"]\n",
        "    blinear1 = weights[\"fc1_biases\" ]\n",
        "\n",
        "    debug=False\n",
        "    for i in range(fconv0.shape[0]):\n",
        "        model.state_dict()['conv0.weight'][i][0]     = torch.from_numpy(fconv0[i]).double()\n",
        "        if debug: print(f\"Weights for Conv0 Filter {i} are {model.state_dict()['conv0.weight'][i][0][0]}   \")\n",
        "\n",
        "    for i in range(fconv1.shape[0]):\n",
        "        for j in range(fconv1.shape[3]):\n",
        "            model.state_dict()['conv1.weight'][i,j,:,:]     = torch.from_numpy(fconv1[i,:,:,j]).double()\n",
        "        if debug: print(f\"Weights for Conv1 Filter {i} are {model.state_dict()['conv1.weight'][i][0][0]}   \")\n",
        "\n",
        "    for i in range(flinear0.shape[1]):\n",
        "        model.state_dict()['linear0.weight'][i]      = torch.from_numpy(flinear0[:,i]).double()\n",
        "        if debug and i<10: print(f\"Weights for Linear Filter {i} are {model.state_dict()['linear0.weight'][i][:5]} \")\n",
        "        \n",
        "    for i in range(blinear0.shape[0]):\n",
        "        model.state_dict()['linear0.bias'][i]        = torch.tensor(blinear0[i]).double()\n",
        "        if debug and i<10: print(f\"Weights for Linear Bias   {i} are {model.state_dict()['linear0.bias'][i]}   \")\n",
        "    \n",
        "    for i in range(flinear1.shape[1]):\n",
        "        model.state_dict()['linear1.weight'][i]      = torch.from_numpy(flinear1[:,i]).double()\n",
        "        if debug and i<10: print(f\"Weights for Linear Filter {i} are {model.state_dict()['linear1.weight'][i][:5]} \")\n",
        "        \n",
        "    for i in range(blinear1.shape[0]):\n",
        "        model.state_dict()['linear1.bias'][i]        = torch.tensor(blinear1[i]).double()\n",
        "        if debug and i<10: print(f\"Weights for Linear Bias   {i} are {model.state_dict()['linear1.bias'][i]}   \")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qnbpVtumU8CF"
      },
      "outputs": [],
      "source": [
        "# Define the Train function\n",
        "def train(epoch):\n",
        "    model.train()\n",
        "    model.double()\n",
        "    copy_weights()\n",
        "    for batch_idx, (image, target) in enumerate(zip(train_images,train_labels)):\n",
        "        optimizer.zero_grad()\n",
        "        \n",
        "        # Prepare Image\n",
        "        img_numpy  = image[:,:,0]\n",
        "        img_tensor = torch.tensor(img_numpy).double()\n",
        "        img_for_model = img_tensor.unsqueeze(0).unsqueeze(0).double()\n",
        "\n",
        "        # Prepare label\n",
        "        label_vector = torch.zeros(1,10)\n",
        "        label_vector[:,target]=1\n",
        "\n",
        "        # Model Output\n",
        "        output = model(img_for_model.double())\n",
        "        \n",
        "        # Calculate Loss\n",
        "        loss = calculate_loss(label_vector, output)\n",
        "        # loss_nlll = calculate_loss_nlll(label_vector, output)\n",
        "        \n",
        "        # Calculate gradients\n",
        "        # scaler.scale(loss).backward()\n",
        "        loss.backward()\n",
        "        \n",
        "        # Update weights\n",
        "        # optimizer.step()\n",
        "        scaler.step(optimizer)\n",
        "        scaler.update()\n",
        "        \n",
        "        grad_bank = {}\n",
        "        for idx, param in enumerate(model.parameters()):\n",
        "            # print(f\"Parameter: {param.name}\")\n",
        "            # print(f\"Parameter: {param.names}\")\n",
        "            # print(f\"Parameter: {param.grad.data}\") \n",
        "            grad_bank[f\"layer_{idx}\"] = param.grad.data\n",
        "        \n",
        "        if model.debug: print(grad_bank['layer_5'])\n",
        "            \n",
        "            \n",
        "        # if batch_idx % log_interval == 0:\n",
        "        print('Train Epoch: {} [{}/{} ({:.0f}%)]\\tLoss: {:.6f}'.format(epoch, batch_idx * len(image), len(train_images),100. * batch_idx / len(train_images), loss.item()),end='\\r')\n",
        "        # train_losses.append(loss.item())\n",
        "        # train_counter.append((batch_idx*64) + ((epoch-1)*len(train_loader.dataset)))\n",
        "        # torch.save(network.state_dict(), '/results/model.pth')\n",
        "        # torch.save(optimizer.state_dict(), '/results/optimizer.pth')\n",
        "        "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "wsiPLY0YU8CF"
      },
      "outputs": [],
      "source": [
        "# Run Training\n",
        "for epoch in range(1, n_epochs + 1):\n",
        "    train(epoch)    "
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "include_colab_link": true,
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3.8.5 ('base')",
      "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.8.5"
    },
    "vscode": {
      "interpreter": {
        "hash": "4fbbc6560802a9753d76c1864ab18e70962bbfcb64c614d1799210197ece64ee"
      }
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}