IDM.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# IDM architecture \n",
    "\n",
    "### Step 1: Load the image dataset "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torchvision \n",
    "import os\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.utils.data import TensorDataset, DataLoader\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "from torch.optim.lr_scheduler import ExponentialLR\n",
    "import gc\n",
    "import numpy as np\n",
    "from torch.utils.tensorboard import SummaryWriter"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Configuration starts "
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Flag to run on GPU or CPU"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "run_on_gpu = True"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Adding tensorboard summary writer for better visualization. Change it before running each time."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "writer = SummaryWriter('runs/experiment-4')"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "No of frames the model will process at a time. For CPU 128 frames work if you have GPU with 12 GB VRAM it will eat up all memory. I tested it with 32 frames and that works."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# No of frame we will process at a time\n",
    "no_of_frames = 32"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Number of data you want to train on"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# -1 means it will use all the test \n",
    "number_of_test_data = -1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "cuda:0\n"
     ]
    }
   ],
   "source": [
    "device = torch.device(\"cuda:0\" if run_on_gpu and torch.cuda.is_available() else \"cpu\")\n",
    "print(device)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Loading the images from the data. \n",
    "First getting the filenames"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "relative_path = 'data'\n",
    "data_dir = os.path.abspath(relative_path)\n",
    "\n",
    "# This below code is only for Anirudha's PC because after restarting Python kernel the abspath gets a bit messed up\n",
    "# Comment it out or change it if you are running on different place\n",
    "# data_dir = 'C:\\\\Users\\\\aniru\\\\Desktop\\\\Code\\\\VPTAirsim\\\\data'\n",
    "filenames = [name for name in os.listdir(data_dir) if os.path.splitext(name)[-1] == '.png']\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Same thing for validation data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "val_relative_path = 'validation'\n",
    "validation_data_dir = os.path.abspath(val_relative_path)\n",
    "\n",
    "# This below code is only for Anirudha's PC because after restarting Python kernel the abspath gets a bit messed up\n",
    "# Comment it out or change it if you are running on different place\n",
    "# validation_data_dir = 'C:\\\\Users\\\\aniru\\\\Desktop\\\\Code\\\\VPTAirsim\\\\validation'\n",
    "validation_filenames = [name for name in os.listdir(validation_data_dir) if os.path.splitext(name)[-1] == '.png']\n",
    "\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Configuration ends "
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Sort the filenames by the number in the string"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_number_from_string(string):\n",
    "    # Extracts the number from a string by filtering out non-digit characters\n",
    "    return int(''.join(filter(str.isdigit, string)))\n",
    "\n",
    "def sort_array_by_number(array):\n",
    "    # Sorts the array based on the number appearing in each string\n",
    "    return sorted(array, key=get_number_from_string)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "sorted_filenames = sort_array_by_number(filenames)\n",
    "\n",
    "if number_of_test_data != -1:\n",
    "    sorted_filenames = sorted_filenames[0:number_of_test_data]"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Same thing for Validation set"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "sorted_val_filenames = sort_array_by_number(validation_filenames)\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Generating the labels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "labels = []\n",
    "\n",
    "for name in sorted_filenames:\n",
    "    \n",
    "    if \"w\" in name:\n",
    "        labels.append(0)\n",
    "    elif \"a\" in name:\n",
    "        labels.append(1)\n",
    "    elif \"d\" in name:\n",
    "        labels.append(2)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "val_labels = []\n",
    "\n",
    "for name in sorted_val_filenames:\n",
    "    \n",
    "    if \"w\" in name:\n",
    "        val_labels.append(0)\n",
    "    elif \"a\" in name:\n",
    "        val_labels.append(1)\n",
    "    elif \"d\" in name:\n",
    "        val_labels.append(2)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Making one hot label"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "labels = torch.tensor(labels)\n",
    "\n",
    "# Perform one-hot encoding\n",
    "one_hot_labels = torch.eye(3)[labels]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "val_labels = torch.tensor(val_labels)\n",
    "\n",
    "# Perform one-hot encoding\n",
    "val_one_hot_labels = torch.eye(3)[val_labels]"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then using the filenames to load images.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Making sure the total data is a multiplication of no of frames \n",
    "dataset_size = len(sorted_filenames) - (len(sorted_filenames) % no_of_frames)\n",
    "\n",
    "\n",
    "dataset = torch.zeros(dataset_size, 3, 128, 128)\n",
    "for i in range(dataset_size):\n",
    " dataset[i] = torchvision.io.read_image(os.path.join(data_dir, sorted_filenames[i]))\n",
    " "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Making sure the total data is a multiplication of no of frames \n",
    "val_dataset_size = len(sorted_val_filenames) - (len(sorted_val_filenames) % no_of_frames)\n",
    "\n",
    "\n",
    "val_dataset = torch.zeros(val_dataset_size, 3, 128, 128)\n",
    "for i in range(val_dataset_size):\n",
    " val_dataset[i] = torchvision.io.read_image(os.path.join(validation_data_dir, sorted_val_filenames[i]))\n",
    " "
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Normalizing the pixel value by dividing it by 255. Now it is between 0 and 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([43808, 3, 128, 128])\n"
     ]
    }
   ],
   "source": [
    "print(dataset.size())\n",
    "dataset = dataset / 255.0\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1280, 3, 128, 128])\n"
     ]
    }
   ],
   "source": [
    "print(val_dataset.size())\n",
    "val_dataset = val_dataset / 255.0"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is to make no_of_frames size chunks of frames"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a TensorDataset with data and their label\n",
    "one_hot_labels = one_hot_labels[0:dataset_size]\n",
    "data = TensorDataset(dataset, one_hot_labels)\n",
    "# Create a DataLoader\n",
    "dataloader = DataLoader(data, batch_size=no_of_frames, shuffle=False)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a TensorDataset with data and their label\n",
    "val_one_hot_labels = val_one_hot_labels[0:val_dataset_size]\n",
    "val_data = TensorDataset(val_dataset, val_one_hot_labels)\n",
    "# Create a DataLoader\n",
    "val_dataloader = DataLoader(val_data, batch_size=no_of_frames, shuffle=False)\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 2 Pass the data through 3D Convolution"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is not a final model. I am writing this to make proper data shape. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Temporal3DConv(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(Temporal3DConv, self).__init__()\n",
    "\n",
    "        # 3 is input channel because of RGB images. \n",
    "        # 128 is the output channel or learnable filters\n",
    "        # Kernel size 5 is temporal kernel width \n",
    "        # (1*1) is spatial kernel width\n",
    "        # 2 Depth padding for initial and end frames\n",
    "        self.conv3d = nn.Conv3d(3, 128, kernel_size=(5, 1, 1), padding=(2,0,0))\n",
    "        self.relu = nn.ReLU()\n",
    "\n",
    "    def forward(self, x):\n",
    "        out = self.conv3d(x)\n",
    "        out = self.relu(out)\n",
    "        return out"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 3 Pass the 3D Convolution layer outcome through ResNet"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is the middle Res Net layer. So the ResNetBlock represents the Deep residual learning for image\n",
    "recognition paper architecture. \n",
    "\n",
    "ResNetBlocksWithPooling represents the Resnet stack mentioned in the VPT paper. We will use three stacks consecuvely then flatten it before passing it to attention layer. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ResNetBlock(nn.Module):\n",
    "    def __init__(self, in_channels, out_channels):\n",
    "        super(ResNetBlock, self).__init__()\n",
    "        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1)\n",
    "        self.relu = nn.ReLU(inplace=True)\n",
    "        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1)\n",
    "\n",
    "        self.shortcut = nn.Sequential()\n",
    "        if in_channels != out_channels:\n",
    "            self.shortcut = nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1)\n",
    "\n",
    "    def forward(self, x):\n",
    "        residual = self.shortcut(x)\n",
    "\n",
    "        out = self.conv1(x)\n",
    "        out = self.relu(out)\n",
    "\n",
    "        out = self.conv2(out)\n",
    "        out += residual\n",
    "        out = self.relu(out)\n",
    "\n",
    "        return out\n",
    "\n",
    "class ResNetBlocksWithPooling(nn.Module):\n",
    "    def __init__(self, in_channels, out_channels):\n",
    "        super(ResNetBlocksWithPooling, self).__init__()\n",
    "        self.conv = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1)\n",
    "        self.pool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)\n",
    "        self.resnet_block1 = ResNetBlock(out_channels, out_channels)\n",
    "        self.resnet_block2 = ResNetBlock(out_channels, out_channels)\n",
    "\n",
    "    def forward(self, x):\n",
    "        out = self.conv(x)\n",
    "        out = self.pool(out)\n",
    "        out = self.resnet_block1(out)\n",
    "        out = self.resnet_block2(out)\n",
    "        return out"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 4 Pass ResNet outcome through Multiheaded Residual Transformer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "class FrameWiseDense(nn.Module):\n",
    "    def __init__(self, in_features, out_features):\n",
    "        super(FrameWiseDense, self).__init__()\n",
    "        self.linear = nn.Linear(in_features, out_features)\n",
    "        self.relu = nn.ReLU()\n",
    "\n",
    "    def forward(self, x):\n",
    "        out = self.linear(x)\n",
    "        out = self.relu(out)\n",
    "        return out\n",
    "\n",
    "class ResidualTransformerBlock(nn.Module):\n",
    "    def __init__(self, embedding_dim, num_heads, dropout=0.1):\n",
    "        super(ResidualTransformerBlock, self).__init__()\n",
    "        self.attention = nn.MultiheadAttention(embedding_dim, num_heads, dropout=dropout)\n",
    "        # self.dropout1 = nn.Dropout(dropout)\n",
    "        self.norm1 = nn.LayerNorm(embedding_dim)\n",
    "        self.dense1 = FrameWiseDense(embedding_dim, 16384)\n",
    "        # self.dropout2 = nn.Dropout(dropout)\n",
    "        self.norm2 = nn.LayerNorm(embedding_dim)\n",
    "        self.dense2 = FrameWiseDense(16384, embedding_dim)\n",
    "\n",
    "    def forward(self, x):\n",
    "        residual = x\n",
    "        out, _ = self.attention(x, x, x)\n",
    "        # out = self.dropout1(out)\n",
    "        out = self.norm1(out + residual)\n",
    "        residual = out\n",
    "        out = self.dense1(out)\n",
    "        # out = self.dropout2(out)\n",
    "        out = self.dense2(out)\n",
    "        out = self.norm2(out + residual)\n",
    "        return out\n",
    "\n",
    "class ActionPredictionModel(nn.Module):\n",
    "    def __init__(self, num_actions):\n",
    "        super(ActionPredictionModel, self).__init__()\n",
    "        # the initial value is 16384 because it is the flattened output dimension for ResNet.\n",
    "        #  For 1 ResNet it is 262144\n",
    "        #  For 2 ResNet it is 65536\n",
    "        # For 3 ResNet it is 16384\n",
    "        self.dense1 = FrameWiseDense(65536, 256) \n",
    "        self.dense2 = FrameWiseDense(256, 4096)  # first 4096 was 256\n",
    "        self.residual_transformer_blocks = nn.Sequential(\n",
    "            ResidualTransformerBlock(embedding_dim=4096, num_heads=32),\n",
    "            # ResidualTransformerBlock(embedding_dim=4096, num_heads=32),\n",
    "            # ResidualTransformerBlock(embedding_dim=4096, num_heads=32),\n",
    "            # ResidualTransformerBlock(embedding_dim=4096, num_heads=32)\n",
    "        )\n",
    "        self.dense3 = FrameWiseDense(4096, 16384)\n",
    "        self.dense4 = FrameWiseDense(16384, 4096)\n",
    "        self.action_head = nn.Linear(4096, num_actions)\n",
    "\n",
    "    def forward(self, x):\n",
    "        out = self.dense1(x)\n",
    "        out = self.dense2(out)\n",
    "        # out = out.permute(1, 0, 2)\n",
    "        out = self.residual_transformer_blocks(out)\n",
    "        # out = out.permute(1, 0, 2)\n",
    "        out = self.dense3(out)\n",
    "        out = self.dense4(out)\n",
    "        # out = out.mean(dim=1)\n",
    "        out = self.action_head(out)\n",
    "        out = F.softmax(out, dim=1)\n",
    "        return out"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 5 Combining everything in a single model "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "class IDM(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(IDM, self).__init__()\n",
    "        self.temporal3DConv = Temporal3DConv()\n",
    "        self.ResNetStack1 = ResNetBlocksWithPooling(128, 64)\n",
    "        self.ResNetStack2 = ResNetBlocksWithPooling(64, 64)\n",
    "        # self.ResNetStack3 = ResNetBlocksWithPooling(64, 64)\n",
    "        self.flattenLayer = nn.Flatten()\n",
    "        self.transformerLayer = ActionPredictionModel(num_actions=3)\n",
    "\n",
    "    def forward(self, input):\n",
    "        # For Conv3D the input format is (batch_size, num_channels, num_frames, height, width)\n",
    "        # So I am using unsqueeze to increase the outer dimension to make batch_size = 1 . \n",
    "        # Then using the permute to make the dimension in proper shape.\n",
    "        out = input.unsqueeze(0).permute(0, 2, 1, 3, 4)\n",
    "        out = self.temporal3DConv(out)\n",
    "        \n",
    "        # To match the expected input shape of the ResNet model, we need to reshape the output tensor. \n",
    "        # First, we use permute to rearrange the dimensions of the tensor, swapping the second and third dimensions. \n",
    "        # Then, we use contiguous to ensure the tensor's memory is laid out contiguously. \n",
    "        # Finally, we use view to reshape the tensor into a 4D tensor with dimensions (batch_size * num_frames, num_channels, height, width).\n",
    "\n",
    "        out = out.permute(0, 2, 1, 3, 4).contiguous().view(1 * no_of_frames, 128, 128, 128)\n",
    "        out = self.ResNetStack1(out)\n",
    "        out = self.ResNetStack2(out)\n",
    "        # out = self.ResNetStack3(out)\n",
    "        out = self.flattenLayer(out)\n",
    "        out = self.transformerLayer(out)\n",
    "        return out\n",
    "\n",
    "    "
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Some codes to check GPU memory status"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "# print(torch.cuda.mem_get_info())\n",
    "# torch.cuda.empty_cache()\n",
    "# torch.cuda.mem_get_info()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 6 Declaring the model"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Loading and printing the model "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# model = IDM().to(device)\n",
    "# print(model)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Checkinng number of trainable parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "# total_params = sum(p.numel() for p in model.parameters() if p.requires_grad)\n",
    "# print(f\"Total Trainable Parameters: {total_params}\")\n",
    "# print(next(model.parameters()).is_cuda)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step 7 Doing Forward pass (NEED TO MODIFY TO SUPPORT FULL TRAINING)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "# output = any\n",
    "\n",
    "# for framechunk in dataloader:\n",
    "#     # Access the batched tensor data\n",
    "#     # Pass the input through the model\n",
    "#     input = framechunk[0].to(device)\n",
    "#     true_label = framechunk[1].to(device)\n",
    "#     output = model(input)\n",
    "#     del input\n",
    "#     torch.cuda.empty_cache()\n",
    "#     gc.collect()\n",
    "#     print(output.shape)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Writing the training loop "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/20, Loss: 1.0616713762283325\n",
      "Validation accuracy =  0.33671875\n",
      "Epoch 2/20, Loss: 1.052653193473816\n",
      "Validation accuracy =  0.3390625\n",
      "Epoch 3/20, Loss: 1.0387365818023682\n",
      "Validation accuracy =  0.33515625\n",
      "Epoch 4/20, Loss: 0.9821683168411255\n",
      "Validation accuracy =  0.38125\n",
      "Epoch 5/20, Loss: 0.9554424285888672\n",
      "Validation accuracy =  0.50390625\n",
      "Epoch 6/20, Loss: 0.850792407989502\n",
      "Validation accuracy =  0.64453125\n",
      "Epoch 7/20, Loss: 0.7347676753997803\n",
      "Validation accuracy =  0.69140625\n",
      "Epoch 8/20, Loss: 0.7381201386451721\n",
      "Validation accuracy =  0.6671875\n",
      "Epoch 9/20, Loss: 0.6167830228805542\n",
      "Validation accuracy =  0.825\n",
      "Epoch 10/20, Loss: 0.592914879322052\n",
      "Validation accuracy =  0.89296875\n",
      "Epoch 11/20, Loss: 0.6028575897216797\n",
      "Validation accuracy =  0.86484375\n",
      "Epoch 12/20, Loss: 0.5695949792861938\n",
      "Validation accuracy =  0.8984375\n",
      "Epoch 13/20, Loss: 0.5740002393722534\n",
      "Validation accuracy =  0.88984375\n",
      "Epoch 14/20, Loss: 0.7851916551589966\n",
      "Validation accuracy =  0.6578125\n",
      "Epoch 15/20, Loss: 0.5658758282661438\n",
      "Validation accuracy =  0.89453125\n",
      "Epoch 16/20, Loss: 0.5664804577827454\n",
      "Validation accuracy =  0.89453125\n",
      "Epoch 17/20, Loss: 0.5661160945892334\n",
      "Validation accuracy =  0.89765625\n",
      "Epoch 18/20, Loss: 0.5671388506889343\n",
      "Validation accuracy =  0.9078125\n",
      "Epoch 19/20, Loss: 0.5661579966545105\n",
      "Validation accuracy =  0.91015625\n",
      "Epoch 20/20, Loss: 0.567590594291687\n",
      "Validation accuracy =  0.9046875\n"
     ]
    }
   ],
   "source": [
    "model = IDM().to(device)\n",
    "\n",
    "# Define your loss function\n",
    "crossLoss = nn.CrossEntropyLoss()\n",
    "\n",
    "# Define your optimizer with weight decay\n",
    "# optimizer = optim.AdamW(model.parameters(), lr=0.003, weight_decay=0.01)\n",
    "\n",
    "# Define the learning rate scheduler\n",
    "# scheduler = ExponentialLR(optimizer, gamma=0.95)\n",
    "\n",
    "# Track the initial model parameters\n",
    "# initial_params = [param.clone().detach() for param in model.parameters()]\n",
    "optimizer = optim.SGD(model.parameters(), lr=0.001)\n",
    "# Training loop\n",
    "num_epochs = 20\n",
    "for epoch in range(num_epochs):\n",
    "    model.train()  # Set the model in training mode\n",
    "    for framechunk in dataloader:\n",
    "        # Access the batched tensor data\n",
    "        # Pass the input through the model\n",
    "        input = framechunk[0].to(device)\n",
    "        true_label = framechunk[1].to(device)\n",
    "        output = model(input)\n",
    "        # print(\"True = \", true_label[0])\n",
    "        # print(\"Output = \", output[0])\n",
    "         # Compute the loss\n",
    "        loss = crossLoss(output, true_label)\n",
    "        optimizer.zero_grad()\n",
    "        # Backward pass\n",
    "        loss.backward()\n",
    "\n",
    "        # Update the model's parameters\n",
    "        optimizer.step()\n",
    "\n",
    "        # del input\n",
    "        # torch.cuda.empty_cache()\n",
    "        # gc.collect()\n",
    "        # print(output.shape)\n",
    "        # Convert softmax predictions to class labels by taking the argmax\n",
    "        # softmax_predictions_labels = np.argmax(output.cpu().detach().numpy(), axis=1)\n",
    "\n",
    "        # Compare the predicted labels with the one-hot encoded labels\n",
    "        # correct_predictions = np.sum(np.equal(softmax_predictions_labels, np.argmax(true_label.cpu().detach().numpy(), axis=1)))\n",
    "\n",
    "        # Calculate accuracy\n",
    "        # accuracy = correct_predictions / float(softmax_predictions_labels.shape[0])\n",
    "        # print(\"Accuracy:\", accuracy)\n",
    "        # print(f\"Epoch {epoch+1}/{num_epochs}, Loss: {loss.item()}\")\n",
    "    \n",
    "    # Check if the model parameters have been updated\n",
    "    # for i, param in enumerate(model.parameters()):\n",
    "    #     if not torch.all(torch.eq(initial_params[i], param)):\n",
    "    #         print(f\"Parameters {i} have been updated.\")\n",
    "\n",
    "    # Update the initial parameters for the next epoch\n",
    "    # initial_params = [param.clone().detach() for param in model.parameters()]\n",
    "\n",
    "    # Update the learning rate\n",
    "    # scheduler.step()\n",
    "    # Validation\n",
    "    model.eval()  # Set the model in evaluation mode\n",
    "    accuracy_list = []\n",
    "    with torch.no_grad():\n",
    "        for framechunk in val_dataloader:\n",
    "            # Access the batched tensor data\n",
    "            # Pass the input through the model\n",
    "            val_input = framechunk[0].to(device)\n",
    "            val_true_label = framechunk[1].to(device)\n",
    "            val_output = model(val_input)\n",
    "\n",
    "            # Convert softmax predictions to class labels by taking the argmax\n",
    "            val_softmax_predictions_labels = np.argmax(val_output.cpu().detach().numpy(), axis=1)\n",
    "\n",
    "            # Compare the predicted labels with the one-hot encoded labels\n",
    "            val_correct_predictions = np.sum(np.equal(val_softmax_predictions_labels, np.argmax(val_true_label.cpu().detach().numpy(), axis=1)))\n",
    "\n",
    "            # Calculate accuracy\n",
    "            val_accuracy = val_correct_predictions / float(val_softmax_predictions_labels.shape[0])\n",
    "            # print(\"val Accuracy:\", val_accuracy)\n",
    "            accuracy_list.append(val_accuracy)\n",
    "\n",
    "        \n",
    "\n",
    "    # Print the loss for this epoch\n",
    "    print(f\"Epoch {epoch+1}/{num_epochs}, Loss: {loss.item()}\")\n",
    "    # Print accuracy on validation set \n",
    "    avg_val_accuracy = sum(accuracy_list)/ len(accuracy_list)\n",
    "    print(\"Validation accuracy = \", avg_val_accuracy)\n",
    "    # Log the loss\n",
    "    writer.add_scalar('Loss/train', loss.item(), epoch)\n",
    "\n",
    "    # Log the accuracy\n",
    "    writer.add_scalar('Accuracy/train', avg_val_accuracy, epoch)\n",
    "writer.close()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This theoritically should clear up the memory but this is not working properly"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [],
   "source": [
    "# del model\n",
    "# torch.cuda.empty_cache()\n",
    "# gc.collect()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.10"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}