Experiment 2 TransNet-MSI-r16.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "f70834c2",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "import os\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0\"  # Select the GPU index\n",
    "import scipy.io\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from torch.nn.parameter import Parameter\n",
    "from torch.nn.init import xavier_uniform_\n",
    "from torch.nn.init import constant_\n",
    "from torch.nn.init import xavier_normal_\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "import math\n",
    "from PIL import Image\n",
    "from collections import OrderedDict\n",
    "import torchvision.transforms as transforms\n",
    "import torchvision.transforms.functional as TF\n",
    "import torchvision.transforms as T\n",
    "import matplotlib.pyplot as plt\n",
    "import time\n",
    "from torch.optim.lr_scheduler import _LRScheduler\n",
    "import warnings\n",
    "import spacy\n",
    "from scipy.io import savemat\n",
    "from scipy import stats\n",
    "import dill as pickle\n",
    "import thop\n",
    "from typing import Optional, Tuple, Any\n",
    "from typing import List, Optional, Tuple\n",
    "from torch_challenge_dataset import DeepVerseChallengeLoaderTaskTwo\n",
    "warnings.filterwarnings(\"ignore\", category=DeprecationWarning)\n",
    "warnings.filterwarnings(\"ignore\", category=UserWarning)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "3fcf443a",
   "metadata": {},
   "outputs": [],
   "source": [
    "Tensor = torch.Tensor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "75a6c22f",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Parameters\n",
    "onoffdict={'GPS': True, 'CAMERAS': True, 'RADAR': True}\n",
    "lr=1e-3\n",
    "num_epochs=100\n",
    "\n",
    "reduction = 16\n",
    "batch_size = 200\n",
    "num_H = 64\n",
    "weight_path=f'models/TransNettask2/cr{reduction}/gps{onoffdict[\"GPS\"]}_cam{onoffdict[\"CAMERAS\"]}_rad{onoffdict[\"RADAR\"]}/'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "8ef315b8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "16"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "reduction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "036636ab",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'models/TransNettask2/cr16/gpsTrue_camTrue_radTrue/'"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "weight_path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "2dae38f0",
   "metadata": {},
   "outputs": [],
   "source": [
    "if not os.path.exists(weight_path):\n",
    "    os.makedirs(weight_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "40a3e63e",
   "metadata": {},
   "outputs": [],
   "source": [
    "device = 'cuda' if torch.cuda.is_available() else 'cpu'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "cce5db81",
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dataset = DeepVerseChallengeLoaderTaskTwo(csv_path = r'./dataset_train.csv')\n",
    "train_loader = DataLoader(train_dataset, batch_size = batch_size, shuffle=True, num_workers=5)\n",
    "test_dataset = DeepVerseChallengeLoaderTaskTwo(csv_path = r'./dataset_validation.csv')\n",
    "test_loader = DataLoader(test_dataset, batch_size = batch_size, shuffle=True, num_workers=5)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c78c6170",
   "metadata": {},
   "source": [
    "# Utils and Models"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "ab275daf",
   "metadata": {},
   "outputs": [],
   "source": [
    "def CSI_abs_reshape(y, csi_std=2.8117975e-06, target_std=1.0):\n",
    "    y = torch.abs(y)\n",
    "    y=(y/csi_std)*target_std\n",
    "    return y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "a8e01a01",
   "metadata": {},
   "outputs": [],
   "source": [
    "def CSI_reshape( y, csi_std=2.5e-06, target_std=1):\n",
    "        ry = torch.real(y)\n",
    "        iy= torch.imag(y)\n",
    "        oy=torch.cat([ry,iy],dim=1)\n",
    "        #scaling\n",
    "        oy=(oy/csi_std)*target_std\n",
    "        return oy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "0b783bf7",
   "metadata": {},
   "outputs": [],
   "source": [
    "def cal_model_parameters(model):\n",
    "    total_param  = []\n",
    "    for p1 in model.parameters():\n",
    "        total_param.append(int(p1.numel()))\n",
    "    return sum(total_param)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "5ee9b3ea",
   "metadata": {},
   "outputs": [],
   "source": [
    "def normalize_image(image):\n",
    "    # Convert image to float tensor\n",
    "    image = image.float()\n",
    "    # Normalize the image\n",
    "    image /= 255.0\n",
    "    # ImageNet mean values # ImageNet standard deviation values\n",
    "    trans=T.Normalize([0.485, 0.456, 0.406],[0.229, 0.224, 0.225]) \n",
    "    image=trans(image)\n",
    "    return image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "2581a147",
   "metadata": {},
   "outputs": [],
   "source": [
    "def left_coordinates_batch(x_cor, y_cor):\n",
    "    y_pix = torch.zeros_like(x_cor)\n",
    "    x_pix = torch.zeros_like(y_cor)\n",
    "\n",
    "    condition1 = y_cor < -4\n",
    "    condition2 = (y_cor >= -4) & (y_cor < -1)\n",
    "    condition3 = (y_cor >= -1) & (y_cor < 1)\n",
    "    condition4 = (y_cor >= 1) & (y_cor < 4)\n",
    "    condition5 = y_cor >= 4\n",
    "\n",
    "    y_pix[condition1] = 100 + (250 - 100) * ((x_cor[condition1] - 80) / (125 - 80))\n",
    "    y_pix[condition2] = 100 + (250 - 100) * ((x_cor[condition2] - 80) / (125 - 80))\n",
    "    y_pix[condition3] = 100 + (250 - 100) * ((x_cor[condition3] - 80) / (125 - 80))\n",
    "    y_pix[condition4] = 100 + (210 - 100) * ((x_cor[condition4] - 80) / (125 - 80))\n",
    "    y_pix[condition5] = 100 + (190 - 100) * ((x_cor[condition5] - 80) / (125 - 80))\n",
    "\n",
    "    x_pix[condition1] = (y_pix[condition1] - 30) / 1.35\n",
    "    x_pix[condition2] = (y_pix[condition2] - 45) / 0.85\n",
    "    x_pix[condition3] = (y_pix[condition3] - 55) / 0.70\n",
    "    x_pix[condition4] = (y_pix[condition4] - 65) / 0.60\n",
    "    x_pix[condition5] = (y_pix[condition5] - 65) / 0.5\n",
    "    return x_pix, y_pix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "c2c4777c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def center_coordinates_batch(x_cor, y_cor):\n",
    "    x_pix = torch.zeros_like(x_cor)\n",
    "    y_pix = torch.zeros_like(y_cor)\n",
    "\n",
    "    condition = y_cor < 0\n",
    "    x_pix[condition] = 256 * ((x_cor[condition] - 119) / (139 - 119))\n",
    "    x_pix[~condition] = 256 * ((x_cor[~condition] - 112) / (146 - 113))\n",
    "    \n",
    "    y_pix = 175 + (100 - 175) * ((y_cor - (-7)) / ((7) - (-7)))\n",
    "    return x_pix, y_pix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "3ce57e49",
   "metadata": {},
   "outputs": [],
   "source": [
    "def right_coordinates_batch(x_cor, y_cor):\n",
    "    y_pix = torch.zeros_like(x_cor)\n",
    "    x_pix = torch.zeros_like(y_cor)\n",
    "\n",
    "    condition1 = y_cor < -4\n",
    "    condition2 = (y_cor >= -4) & (y_cor < -1)\n",
    "    condition3 = (y_cor >= -1) & (y_cor < 1)\n",
    "    condition4 = (y_cor >= 1) & (y_cor < 4)\n",
    "    condition5 = y_cor >= 4\n",
    "\n",
    "    y_pix[condition1] = 250 + (100 - 250) * ((x_cor[condition1] - 125) / (200 - 125))\n",
    "    y_pix[condition2] = 250 + (100 - 250) * ((x_cor[condition2] - 125) / (200 - 125))\n",
    "    y_pix[condition3] = 250 + (100 - 250) * ((x_cor[condition3] - 125) / (200 - 125))\n",
    "    y_pix[condition4] = 210 + (100 - 210) * ((x_cor[condition4] - 125) / (200 - 125))\n",
    "    y_pix[condition5] = 190 + (100 - 190) * ((x_cor[condition5] - 125) / (200 - 125))\n",
    "\n",
    "    x_pix[condition1] = -(y_pix[condition1] - 370) / 1.25\n",
    "    x_pix[condition2] = -(y_pix[condition2] - 285) / 0.87\n",
    "    x_pix[condition3] = -(y_pix[condition3] - 250) / 0.73\n",
    "    x_pix[condition4] = -(y_pix[condition4] - 210) / 0.55\n",
    "    x_pix[condition5] = -(y_pix[condition5] - 190) / 0.45\n",
    "    return x_pix, y_pix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "3c98e4f6",
   "metadata": {},
   "outputs": [],
   "source": [
    "def center_image_batch(images, center_x, center_y, output_size, bounded=False):\n",
    "    batch_size = images.size(0)\n",
    "    \n",
    "    if bounded == 'left':\n",
    "        top = torch.clamp(center_y - output_size[0] // 2, 0, None)\n",
    "        left = torch.clamp(center_x - output_size[1] // 2, 0, None)\n",
    "    elif bounded == 'right':\n",
    "        bottom = center_y + output_size[0] // 2\n",
    "        right = torch.clamp(center_x + output_size[1] // 2, None, 250)\n",
    "        top = torch.clamp(bottom - output_size[0], 0, None)\n",
    "        left = right - output_size[1]\n",
    "    else:\n",
    "        top = center_y - output_size[0] // 2\n",
    "        left = center_x - output_size[1] // 2\n",
    "\n",
    "    resize_transform = transforms.Resize((output_size))\n",
    "    cropped_images = [TF.crop(image, int(top[i].item()), int(left[i].item()), output_size[0], output_size[1]) \n",
    "                        for i, image in enumerate(images)]\n",
    "    cropped_images = torch.stack([resize_transform(image) for image in cropped_images])\n",
    "    return cropped_images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "a93af7a9",
   "metadata": {},
   "outputs": [],
   "source": [
    "def process_imgs(gps, img_1, img_2, img_3, crop_size = (150,150)):\n",
    "    x_cor = gps[:,0]\n",
    "    y_cor = gps[:,1]\n",
    "\n",
    "    x_pix,y_pix = left_coordinates_batch(x_cor, y_cor)\n",
    "    img_1 = center_image_batch(img_1, x_pix.to(torch.int), y_pix.to(torch.int), crop_size, 'left')\n",
    "\n",
    "    x_pix,y_pix = center_coordinates_batch(x_cor, y_cor)\n",
    "    img_2 = center_image_batch(img_2, x_pix.to(torch.int), y_pix.to(torch.int), crop_size)\n",
    "\n",
    "    x_pix,y_pix = right_coordinates_batch(x_cor, y_cor)\n",
    "    img_3 = center_image_batch(img_3, x_pix.to(torch.int), y_pix.to(torch.int), crop_size, 'right')\n",
    "\n",
    "    return img_1, img_2, img_3"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a516ccee",
   "metadata": {},
   "source": [
    "Scheduler"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "25de8041",
   "metadata": {},
   "outputs": [],
   "source": [
    "class gpsdata(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.gps_fc = nn.Linear(2, 16)\n",
    "        self.gps_relu = nn.ReLU()\n",
    "\n",
    "    def forward(self, gps):  \n",
    "        gps = gps.to(torch.float32)\n",
    "\n",
    "        x, y = gps[:,0], gps[:,1]\n",
    "        x_normd = (x - torch.min(x)) / (torch.max(x) - torch.min(x))\n",
    "        y_normd = (y - torch.min(y)) / (torch.max(y) - torch.min(y))\n",
    "        gps_normd = torch.stack([x_normd,y_normd],dim=1)\n",
    "\n",
    "        gps_out = self.gps_fc(gps_normd)  \n",
    "        gps_out = self.gps_relu(gps_out)\n",
    "        return gps_out"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "38a26d3e",
   "metadata": {},
   "source": [
    "Radar Data Processing Layer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "6e65f75c",
   "metadata": {},
   "outputs": [],
   "source": [
    "class radardata(nn.Module):\n",
    "    \n",
    "    def __init__(self):\n",
    "        super(radardata, self).__init__()\n",
    "        self.dropout = nn.Dropout(p=0.5)\n",
    "        self.conv1 = nn.Conv2d(1, 6, kernel_size=3, stride=2, padding=1)\n",
    "        self.pool1 = nn.AvgPool2d(kernel_size=2, stride=2)\n",
    "        self.conv2 = nn.Conv2d(6, 1, kernel_size=3, stride=2, padding=1)\n",
    "        self.pool2 = nn.AvgPool2d(kernel_size=2, stride=2)\n",
    "        self.lr1=nn.LeakyReLU(negative_slope=0.3, inplace=True)\n",
    "        self.encoder_fc = nn.Linear(256,16)\n",
    "                \n",
    "    def forward(self, x):\n",
    "        x=x.view(-1,1,512,128)\n",
    "        x = (x - 5.1838e-06) / (28.0494 - 5.1838e-06)\n",
    "        out = self.pool1(self.dropout(self.conv1(x)))\n",
    "        out = self.pool2(self.dropout(self.conv2(out))).view(x.size(0), -1)\n",
    "        out = self.dropout(self.encoder_fc(out))\n",
    "        out = self.lr1(out)\n",
    "        return out "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9004a71e",
   "metadata": {},
   "source": [
    "Camera Data Processing Layers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "5dce3f9e",
   "metadata": {},
   "outputs": [],
   "source": [
    "class cameradata(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.conv1 = nn.Conv2d(3, 6, kernel_size=3, stride=2, padding=1)\n",
    "        self.pool1 = nn.AvgPool2d(kernel_size=2, stride=2)\n",
    "        self.conv2 = nn.Conv2d(6, 1, kernel_size=3, stride=2, padding=1)\n",
    "        self.pool2 = nn.AvgPool2d(kernel_size=2, stride=2)\n",
    "        self.lr1=nn.LeakyReLU(negative_slope=0.3, inplace=True)\n",
    "        self.encoder = nn.Linear(1*81,16)\n",
    "        self.dropout = nn.Dropout(0.5)\n",
    "\n",
    "    def forward(self, cam):  \n",
    "        cam = normalize_image(cam).to(torch.float32)\n",
    "        out = self.pool1(self.dropout(self.conv1(cam)))\n",
    "        out = self.pool2(self.dropout(self.conv2(out)))\n",
    "        out = self.lr1(out).view(-1,1*81)\n",
    "        out = self.dropout(self.encoder(out))\n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "c21d6cac",
   "metadata": {},
   "outputs": [],
   "source": [
    "def scale_dot_attention(\n",
    "       q:Tensor,\n",
    "       k:Tensor,\n",
    "       v:Tensor,\n",
    "       dropout_p:float = 0.0,\n",
    "       attn_mask : Optional[Tensor] = None,\n",
    ")-> Tuple[Tensor,Tensor]:\n",
    "    \n",
    "    _,_,E = q.shape\n",
    "    q = q / math.sqrt(E)\n",
    "    attn = torch.bmm(q,k.transpose(-2,-1))\n",
    "    if attn_mask is not None:\n",
    "        attn = attn + attn_mask\n",
    "    attn = F.softmax(attn,dim =-1)\n",
    "    if dropout_p:\n",
    "        attn = F.dropout(attn,p = dropout_p)\n",
    "    out = torch.bmm(attn,v)\n",
    "\n",
    "    return out,attn\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "89e3110f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def multi_head_attention_forward(\n",
    "        query: Tensor,\n",
    "        key: Tensor,\n",
    "        value: Tensor,\n",
    "        num_heads: int,\n",
    "        in_proj_weight: Tensor,\n",
    "        in_proj_bias: Optional[Tensor],\n",
    "        dropout_p: float,\n",
    "        out_proj_weight: Tensor,\n",
    "        out_proj_bias: Optional[Tensor],\n",
    "        training: bool = True,\n",
    "        key_padding_mask: Optional[Tensor] = None,\n",
    "        need_weights: bool = True,\n",
    "        attn_mask: Optional[Tensor] = None,\n",
    "        use_separate_proj_weight=None,\n",
    "        q_proj_weight: Optional[Tensor] = None,\n",
    "        k_proj_weight: Optional[Tensor] = None,\n",
    "        v_proj_weight: Optional[Tensor] = None,\n",
    ") -> Tuple[Tensor, Optional[Tensor]]:\n",
    "    \n",
    "    tgt_len, bsz, embed_dim = query.shape\n",
    "    src_len, _, _ = key.shape\n",
    "    head_dim = embed_dim // num_heads\n",
    "    q, k, v = _in_projection_packed(query, key, value, in_proj_weight, in_proj_bias)\n",
    "    \n",
    "    if attn_mask is not None:\n",
    "        if attn_mask.dtype == torch.uint8:\n",
    "            warnings.warn(\"Byte tensor for attn_mask in nn.MultiheadAttention is deprecated. Use bool tensor instead.\")\n",
    "            attn_mask = attn_mask.to(torch.bool)\n",
    "        else:\n",
    "            assert attn_mask.is_floating_point() or attn_mask.dtype == torch.bool, \\\n",
    "                f\"Only float, byte, and bool types are supported for attn_mask, not {attn_mask.dtype}\"\n",
    "\n",
    "        if attn_mask.dim() == 2:\n",
    "            correct_2d_size = (tgt_len, src_len)\n",
    "            if attn_mask.shape != correct_2d_size:\n",
    "                raise RuntimeError(\n",
    "                    f\"The shape of the 2D attn_mask is {attn_mask.shape}, but should be {correct_2d_size}.\")\n",
    "            attn_mask = attn_mask.unsqueeze(0)\n",
    "        elif attn_mask.dim() == 3:\n",
    "            correct_3d_size = (bsz * num_heads, tgt_len, src_len)\n",
    "            if attn_mask.shape != correct_3d_size:\n",
    "                raise RuntimeError(\n",
    "                    f\"The shape of the 3D attn_mask is {attn_mask.shape}, but should be {correct_3d_size}.\")\n",
    "        else:\n",
    "            raise RuntimeError(f\"attn_mask's dimension {attn_mask.dim()} is not supported\")\n",
    "\n",
    "    if key_padding_mask is not None and key_padding_mask.dtype == torch.uint8:\n",
    "        warnings.warn(\n",
    "            \"Byte tensor for key_padding_mask in nn.MultiheadAttention is deprecated. Use bool tensor instead.\")\n",
    "        key_padding_mask = key_padding_mask.to(torch.bool)\n",
    "\n",
    "\n",
    "    q = q.contiguous().view(tgt_len, bsz * num_heads, head_dim).transpose(0, 1)\n",
    "    k = k.contiguous().view(-1, bsz * num_heads, head_dim).transpose(0, 1)\n",
    "    v = v.contiguous().view(-1, bsz * num_heads, head_dim).transpose(0, 1)\n",
    "    \n",
    "    if key_padding_mask is not None:\n",
    "        assert key_padding_mask.shape == (bsz, src_len), \\\n",
    "            f\"expecting key_padding_mask shape of {(bsz, src_len)}, but got {key_padding_mask.shape}\"\n",
    "        key_padding_mask = key_padding_mask.view(bsz, 1, 1, src_len). \\\n",
    "            expand(-1, num_heads, -1, -1).reshape(bsz * num_heads, 1, src_len)\n",
    "        if attn_mask is None:\n",
    "            attn_mask = key_padding_mask\n",
    "        elif attn_mask.dtype == torch.bool:\n",
    "            attn_mask = attn_mask.logical_or(key_padding_mask)\n",
    "        else:\n",
    "            attn_mask = attn_mask.masked_fill(key_padding_mask, float(\"-inf\"))\n",
    "  \n",
    "    if attn_mask is not None and attn_mask.dtype == torch.bool:\n",
    "        new_attn_mask = torch.zeros_like(attn_mask, dtype=torch.float)\n",
    "        new_attn_mask.masked_fill_(attn_mask, float(\"-inf\"))\n",
    "        attn_mask = new_attn_mask\n",
    "\n",
    "    \n",
    "    if not training:\n",
    "        dropout_p = 0.0\n",
    "    attn_output, attn_output_weights = scale_dot_attention(q, k, v, attn_mask, dropout_p)\n",
    "    attn_output = attn_output.transpose(0, 1).contiguous().view(tgt_len, bsz, embed_dim)\n",
    "    attn_output = nn.functional.linear(attn_output, out_proj_weight, out_proj_bias)\n",
    "    if need_weights:\n",
    "        # average attention weights over heads\n",
    "        attn_output_weights = attn_output_weights.view(bsz, num_heads, tgt_len, src_len)\n",
    "        return attn_output, attn_output_weights.sum(dim=1) / num_heads\n",
    "    else:\n",
    "        return attn_output, None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "1518927c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def _in_projection_packed(\n",
    "    q: Tensor,\n",
    "    k: Tensor,\n",
    "    v: Tensor,\n",
    "    w: Tensor,\n",
    "    b: Optional[Tensor] = None,\n",
    ") -> List[Tensor]:\n",
    "    E = q.size(-1)\n",
    "    if k is v:\n",
    "        if q is k:\n",
    "            return F.linear(q, w, b).chunk(3, dim=-1)\n",
    "        else:\n",
    "            w_q, w_kv = w.split([E, E * 2])\n",
    "            if b is None:\n",
    "                b_q = b_kv = None\n",
    "            else:\n",
    "                b_q, b_kv = b.split([E, E * 2])\n",
    "            return (F.linear(q, w_q, b_q),) + F.linear(k, w_kv, b_kv).chunk(2, dim=-1)\n",
    "    else:\n",
    "        w_q, w_k, w_v = w.chunk(3)\n",
    "        if b is None:\n",
    "            b_q = b_k = b_v = None\n",
    "        else:\n",
    "            b_q, b_k, b_v = b.chunk(3)\n",
    "        return F.linear(q, w_q, b_q), F.linear(k, w_k, b_k), F.linear(v, w_v, b_v)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "bcdf9016",
   "metadata": {},
   "outputs": [],
   "source": [
    "class MultiheadAttention(nn.Module):\n",
    "\n",
    "    def __init__(self, embed_dim, num_heads, dropout=0., bias=True,\n",
    "                 kdim=None, vdim=None, batch_first=False) -> None:\n",
    "        # factory_kwargs = {'device': device, 'dtype': dtype}\n",
    "        super(MultiheadAttention, self).__init__()\n",
    "        self.embed_dim = embed_dim\n",
    "        self.kdim = kdim if kdim is not None else embed_dim\n",
    "        self.vdim = vdim if vdim is not None else embed_dim\n",
    "        self._qkv_same_embed_dim = self.kdim == self.embed_dim and self.vdim == self.embed_dim\n",
    "\n",
    "        self.num_heads = num_heads\n",
    "        self.dropout = dropout\n",
    "        self.batch_first = batch_first\n",
    "        self.head_dim = embed_dim // num_heads\n",
    "        assert self.head_dim * num_heads == self.embed_dim, \"embed_dim must be divisible by num_heads\"\n",
    "\n",
    "        if self._qkv_same_embed_dim is False:\n",
    "            self.q_proj_weight = Parameter(torch.empty((embed_dim, embed_dim)))\n",
    "            self.k_proj_weight = Parameter(torch.empty((embed_dim, self.kdim)))\n",
    "            self.v_proj_weight = Parameter(torch.empty((embed_dim, self.vdim)))\n",
    "            self.register_parameter('in_proj_weight', None)\n",
    "        else:\n",
    "            self.in_proj_weight = Parameter(torch.empty((3 * embed_dim,embed_dim)))\n",
    "            self.register_parameter('q_proj_weight', None)\n",
    "            self.register_parameter('k_proj_weight', None)\n",
    "            self.register_parameter('v_proj_weight', None)\n",
    "\n",
    "        if bias:\n",
    "            self.in_proj_bias = Parameter(torch.empty(3 * embed_dim))\n",
    "        else:\n",
    "            self.register_parameter('in_proj_bias', None)\n",
    "        self.out_proj = nn.Linear(embed_dim, embed_dim, bias=bias)\n",
    "\n",
    "        self._reset_parameters()\n",
    "        \n",
    "    def _reset_parameters(self):\n",
    "        if self._qkv_same_embed_dim:\n",
    "            xavier_uniform_(self.in_proj_weight)\n",
    "        else:\n",
    "            xavier_uniform_(self.q_proj_weight)\n",
    "            xavier_uniform_(self.k_proj_weight)\n",
    "            xavier_uniform_(self.v_proj_weight)\n",
    "\n",
    "        if self.in_proj_bias is not None:\n",
    "            constant_(self.in_proj_bias, 0.)\n",
    "            constant_(self.out_proj.bias, 0.)\n",
    "\n",
    "\n",
    "\n",
    "    def forward(self, query: Tensor, key: Tensor, value: Tensor, key_padding_mask: Optional[Tensor] = None,\n",
    "                need_weights: bool = True, attn_mask: Optional[Tensor] = None) -> Tuple[Tensor, Optional[Tensor]]:\n",
    "        if self.batch_first:\n",
    "            query, key, value = [x.transpose(1, 0) for x in (query, key, value)]\n",
    "\n",
    "        if not self._qkv_same_embed_dim:\n",
    "            attn_output, attn_output_weights = multi_head_attention_forward(\n",
    "                query, key, value, self.num_heads,\n",
    "                self.in_proj_weight, self.in_proj_bias,\n",
    "                self.dropout, self.out_proj.weight, self.out_proj.bias,\n",
    "                training=self.training,\n",
    "                key_padding_mask=key_padding_mask, need_weights=need_weights,\n",
    "                attn_mask=attn_mask, use_separate_proj_weight=True,\n",
    "                q_proj_weight=self.q_proj_weight, k_proj_weight=self.k_proj_weight,\n",
    "                v_proj_weight=self.v_proj_weight)\n",
    "        else:\n",
    "            attn_output, attn_output_weights = multi_head_attention_forward(\n",
    "                query, key, value, self.num_heads,\n",
    "                self.in_proj_weight, self.in_proj_bias,\n",
    "                self.dropout, self.out_proj.weight, self.out_proj.bias,\n",
    "                training=self.training,\n",
    "                key_padding_mask=key_padding_mask, need_weights=need_weights,\n",
    "                attn_mask=attn_mask)\n",
    "            \n",
    "        if self.batch_first:\n",
    "            return attn_output.transpose(1, 0), attn_output_weights\n",
    "        else:\n",
    "            return attn_output, attn_output_weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "033e780f",
   "metadata": {},
   "outputs": [],
   "source": [
    "class TransformerEncoderLayer(nn.Module):\n",
    "\n",
    "    def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1, activation=F.relu,\n",
    "                 layer_norm_eps=1e-5, batch_first=False) -> None:\n",
    "        super(TransformerEncoderLayer, self).__init__()\n",
    "        self.self_attn = MultiheadAttention(d_model,nhead,\n",
    "                                            dropout=dropout, batch_first=batch_first)\n",
    "\n",
    "        self.linear1 = nn.Linear(d_model, dim_feedforward)\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "        self.linear2 = nn.Linear(dim_feedforward, d_model)\n",
    "\n",
    "        self.norm1 = nn.LayerNorm(d_model, eps=layer_norm_eps)\n",
    "        self.norm2 = nn.LayerNorm(d_model, eps=layer_norm_eps)\n",
    "        self.dropout1 = nn.Dropout(dropout)\n",
    "        self.dropout2 = nn.Dropout(dropout)\n",
    "        self.activation = activation\n",
    "\n",
    "    def forward(self, src: Tensor, src_mask: Optional[Tensor] = None,\n",
    "                src_key_padding_mask: Optional[Tensor] = None) -> Tensor:\n",
    "        src2 = self.self_attn(src, src, src, attn_mask=src_mask,\n",
    "                              key_padding_mask=src_key_padding_mask)[0]\n",
    "        src = src + self.dropout1(src2)\n",
    "        src = self.norm1(src)\n",
    "        src2 = self.linear2(self.dropout(self.activation(self.linear1(src))))\n",
    "        src = src + self.dropout(src2)\n",
    "        src = self.norm2(src)\n",
    "        return src"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "a419510b",
   "metadata": {},
   "outputs": [],
   "source": [
    "class TransformerEncoder(nn.Module):\n",
    "\n",
    "    def __init__(self, encoder_layer, num_layers, norm=None):\n",
    "        super(TransformerEncoder, self).__init__()\n",
    "        self.layer = encoder_layer\n",
    "        self.num_layers = num_layers\n",
    "        self.norm = norm\n",
    "\n",
    "    def forward(self, src: Tensor, mask: Optional[Tensor] = None,\n",
    "                src_key_padding_mask: Optional[Tensor] = None) -> Tensor:\n",
    "        output = src\n",
    "        for _ in range(self.num_layers):\n",
    "            output = self.layer(output, src_mask=mask, src_key_padding_mask=src_key_padding_mask)\n",
    "\n",
    "        if self.norm is not None:\n",
    "            output = self.norm(output)\n",
    "\n",
    "        return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "ec5e8c31",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Decoder Layer:\n",
    "class TransformerDecoderLayer(nn.Module):\n",
    "    def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1, activation=F.relu,\n",
    "                 layer_norm_eps=1e-5, batch_first=False) -> None:\n",
    "        super(TransformerDecoderLayer, self).__init__()\n",
    "        self.self_attn = MultiheadAttention(d_model,nhead,\n",
    "                                            dropout=dropout, batch_first=batch_first)\n",
    "        self.multihead_attn = MultiheadAttention(d_model,nhead,dropout=dropout, batch_first=batch_first)\n",
    "\n",
    "        self.linear1 = nn.Linear(d_model, dim_feedforward)\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "        self.linear2 = nn.Linear(dim_feedforward, d_model)\n",
    "\n",
    "        self.norm1 = nn.LayerNorm(d_model, eps=layer_norm_eps)\n",
    "        self.norm2 = nn.LayerNorm(d_model, eps=layer_norm_eps)\n",
    "        self.norm3 = nn.LayerNorm(d_model, eps=layer_norm_eps)\n",
    "        self.dropout1 = nn.Dropout(dropout)\n",
    "        self.dropout2 = nn.Dropout(dropout)\n",
    "        self.dropout3 = nn.Dropout(dropout)\n",
    "\n",
    "        self.activation = activation\n",
    "\n",
    "    def forward(self, tgt: Tensor, memory: Tensor, tgt_mask: Optional[Tensor] = None,\n",
    "                memory_mask: Optional[Tensor] = None, tgt_key_padding_mask: Optional[Tensor] = None,\n",
    "                memory_key_padding_mask: Optional[Tensor] = None) -> Tensor:\n",
    "        tgt2 = self.self_attn(tgt, tgt, tgt, attn_mask=tgt_mask,\n",
    "                              key_padding_mask=tgt_key_padding_mask)[0]\n",
    "        tgt = tgt + self.dropout1(tgt2)\n",
    "        tgt = self.norm1(tgt)\n",
    "        tgt2 = self.multihead_attn(tgt, memory, memory, attn_mask=memory_mask,\n",
    "                                   key_padding_mask=memory_key_padding_mask)[0]\n",
    "        tgt = tgt + self.dropout2(tgt2)\n",
    "        tgt = self.norm2(tgt)\n",
    "        tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt))))\n",
    "        tgt = tgt + self.dropout3(tgt2)\n",
    "        tgt = self.norm3(tgt)\n",
    "        return tgt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "9146d0b9",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Decoder\n",
    "class TransformerDecoder(nn.Module):\n",
    "    \n",
    "    def __init__(self, decoder_layer, num_layers, norm=None):\n",
    "        super(TransformerDecoder, self).__init__()\n",
    "        self.layer = decoder_layer\n",
    "        self.num_layers = num_layers\n",
    "        self.norm = norm\n",
    "\n",
    "    def forward(self, tgt: Tensor, memory: Tensor, tgt_mask: Optional[Tensor] = None,\n",
    "                memory_mask: Optional[Tensor] = None, tgt_key_padding_mask: Optional[Tensor] = None,\n",
    "                memory_key_padding_mask: Optional[Tensor] = None) -> Tensor:\n",
    "        output = tgt\n",
    "        for _ in range(self.num_layers):\n",
    "            output = self.layer(output, memory, tgt_mask=tgt_mask,\n",
    "                                memory_mask=memory_mask,\n",
    "                                tgt_key_padding_mask=tgt_key_padding_mask,\n",
    "                                memory_key_padding_mask=memory_key_padding_mask)\n",
    "        if self.norm is not None:\n",
    "            output = self.norm(output)\n",
    "\n",
    "        return output\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8f9156eb",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8a189d30",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f625ab39",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "f2f70d47",
   "metadata": {},
   "outputs": [],
   "source": [
    "class task2Encoder(nn.Module):\n",
    "    \n",
    "    def __init__(self,  d_model: int = 64, nhead: int = 8, num_encoder_layers: int = 6,\n",
    "                 num_decoder_layers: int = 6, dim_feedforward: int = 2048, dropout: float = 0.1,\n",
    "                 activation = F.relu, custom_encoder: Optional[Any] = None, custom_decoder: Optional[Any] = None,\n",
    "                 layer_norm_eps: float = 1e-5, batch_first: bool = False,  reduction=64) -> None:\n",
    "        super(task2Encoder, self).__init__()\n",
    "        self.total_size =8192\n",
    "       \n",
    "        if custom_encoder is not None:\n",
    "            self.encoder = custom_encoder\n",
    "        else:\n",
    "            encoder_layer = TransformerEncoderLayer(d_model, nhead, dim_feedforward, dropout,\n",
    "                                                    activation, layer_norm_eps, batch_first)\n",
    "            encoder_norm = nn.LayerNorm(d_model, eps=layer_norm_eps)\n",
    "            self.encoder = TransformerEncoder(encoder_layer, num_encoder_layers)\n",
    "       \n",
    "        self.d_model = d_model\n",
    "        \n",
    "        assert not (self.total_size % self.d_model), 'd_model needs to be divisible by the size of the entire csi matrix (2048)'\n",
    "        self.feature_shape = (self.total_size//(2*self.d_model), self.d_model)\n",
    "\n",
    "        self.nhead = nhead\n",
    "\n",
    "        self.batch_first = batch_first\n",
    "        self.fc_encoder = nn.Linear(self.total_size//2,self.total_size//reduction)\n",
    "        #self.fc_decoder = nn.Linear(self.total_size//reduction,self.total_size)\n",
    "        self._reset_parameters()\n",
    "\n",
    "                \n",
    "        \n",
    "    def forward(self, src: Tensor, tgt: Optional[Tensor]=None, src_mask: Optional[Tensor] = None,\n",
    "                    tgt_mask: Optional[Tensor] = None,\n",
    "                    memory_mask: Optional[Tensor] = None, src_key_padding_mask: Optional[Tensor] = None,\n",
    "                    tgt_key_padding_mask: Optional[Tensor] = None,\n",
    "                    memory_key_padding_mask: Optional[Tensor] = None) -> Tensor:\n",
    "        \n",
    "        memory = self.encoder(src.view(-1, self.feature_shape[0], self.feature_shape[1]), mask=src_mask, src_key_padding_mask=src_key_padding_mask)\n",
    "        memory_encoder = self.fc_encoder(memory.view(memory.shape[0],-1))\n",
    "        \n",
    "        return memory_encoder\n",
    "    def generate_square_subsequent_mask(self, sz: int) -> Tensor:\n",
    "        \n",
    "        mask = (torch.triu(torch.ones(sz, sz)) == 1).transpose(0, 1)\n",
    "\n",
    "        mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0))\n",
    "        return mask\n",
    "\n",
    "    def _reset_parameters(self):\n",
    "   \n",
    "        for p in self.parameters():\n",
    "            if p.dim() > 1:\n",
    "                xavier_uniform_(p)  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "bf1fbe1b",
   "metadata": {},
   "outputs": [],
   "source": [
    "class task2Decoder(nn.Module):\n",
    "    \n",
    "    def __init__(self,  d_model: int = 64, nhead: int = 8, num_encoder_layers: int = 6,\n",
    "                 num_decoder_layers: int = 6, dim_feedforward: int = 2048, dropout: float = 0.1,\n",
    "                 activation = F.relu, custom_encoder: Optional[Any] = None, custom_decoder: Optional[Any] = None,\n",
    "                 layer_norm_eps: float = 1e-5, batch_first: bool = False,  reduction=64) -> None:\n",
    "        super(task2Decoder, self).__init__()\n",
    "        self.total_size = 8192\n",
    "        w, h =64, 64\n",
    "        if custom_decoder is not None:\n",
    "            self.decoder = custom_decoder\n",
    "        else:\n",
    "            decoder_layer = TransformerDecoderLayer(d_model, nhead, dim_feedforward, dropout,\n",
    "                                                    activation, layer_norm_eps, batch_first)\n",
    "            decoder_norm = nn.LayerNorm(d_model, eps=layer_norm_eps)\n",
    "            self.decoder = TransformerDecoder(decoder_layer, num_decoder_layers, decoder_norm)\n",
    "        \n",
    "        self.d_model = d_model\n",
    "        \n",
    "        assert not (self.total_size % self.d_model), 'd_model needs to be divisible by the size of the entire csi matrix (2048)'\n",
    "        self.feature_shape = (self.total_size//self.d_model, self.d_model)\n",
    "\n",
    "        self.nhead = nhead\n",
    "\n",
    "        self.batch_first = batch_first\n",
    "        #self.fc_encoder = nn.Linear(self.total_size,self.total_size//reduction)\n",
    "        self.fc_decoder = nn.Linear(self.total_size//reduction,self.total_size)\n",
    "        self._reset_parameters()\n",
    "        \n",
    "        self.decoder_fc2 = nn.Linear(self.total_size, self.total_size//2)\n",
    "        self.sig2 = nn.Sigmoid()\n",
    "        \n",
    "    \n",
    "    def forward(self, memory_encoder, tgt: Optional[Tensor]=None, src_mask: Optional[Tensor] = None,\n",
    "                    tgt_mask: Optional[Tensor] = None,\n",
    "                    memory_mask: Optional[Tensor] = None, src_key_padding_mask: Optional[Tensor] = None,\n",
    "                    tgt_key_padding_mask: Optional[Tensor] = None,\n",
    "                    memory_key_padding_mask: Optional[Tensor] = None) -> Tensor:\n",
    "        bs = memory_encoder.size(0)\n",
    "        \n",
    "        # Generate final output\n",
    "        memory_decoder = self.fc_decoder(memory_encoder).view(-1, self.feature_shape[0], self.feature_shape[1])\n",
    "        output = self.decoder(memory_decoder, memory_decoder, tgt_mask=tgt_mask, memory_mask=memory_mask,\n",
    "                              tgt_key_padding_mask=tgt_key_padding_mask,\n",
    "                              memory_key_padding_mask=memory_key_padding_mask)\n",
    "        \n",
    "        output = self.sig2(self.decoder_fc2(output.view(bs, -1)))\n",
    "        \n",
    "        return output.view(bs, -1, 64, 64)\n",
    "    \n",
    "    def generate_square_subsequent_mask(self, sz: int) -> Tensor:\n",
    "        \n",
    "        mask = (torch.triu(torch.ones(sz, sz)) == 1).transpose(0, 1)\n",
    "\n",
    "        mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0))\n",
    "        return mask\n",
    "\n",
    "    def _reset_parameters(self):\n",
    "   \n",
    "        for p in self.parameters():\n",
    "            if p.dim() > 1:\n",
    "                xavier_uniform_(p)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "46df19e4",
   "metadata": {},
   "outputs": [],
   "source": [
    "#complete task 2 model including encoder, decoder and channel\n",
    "class task2model(nn.Module):\n",
    "    def __init__(self, reduction=16):\n",
    "        super().__init__()\n",
    "        \n",
    "        self.en=task2Encoder(d_model= 32, num_encoder_layers=2, num_decoder_layers=2, nhead=2, reduction =reduction, dropout= 0.)\n",
    "        \n",
    "        self.de=task2Decoder(d_model=32, num_encoder_layers=2, num_decoder_layers=2, nhead=2, reduction =reduction, dropout= 0.)\n",
    "        \n",
    "    \n",
    "   \n",
    "    def forward(self, Hin, device, is_training): \n",
    "        \n",
    "        #Encoder\n",
    "        Hencoded=self.en(Hin)\n",
    "        \n",
    "        \n",
    "        #Decoder   \n",
    "        Hdecoded=self.de(Hencoded)\n",
    "        \n",
    "\n",
    "        return Hdecoded"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f7ea24f5",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "7dd4bac1",
   "metadata": {},
   "outputs": [],
   "source": [
    "#loading weights of baseline model and task 1\n",
    "onoffdictb={'GPS': False, 'CAMERAS': False, 'RADAR': False} #baseline dictionary\n",
    "weight_pathb=f'models/TransNettask2/cr{reduction}/gps{onoffdictb[\"GPS\"]}_cam{onoffdictb[\"CAMERAS\"]}_rad{onoffdictb[\"RADAR\"]}/'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "69a76c16",
   "metadata": {},
   "outputs": [],
   "source": [
    "task1_weight_path=f'models/task1/gps{onoffdict[\"GPS\"]}_cam{onoffdict[\"CAMERAS\"]}_rad{onoffdict[\"RADAR\"]}/'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "14888412",
   "metadata": {},
   "outputs": [],
   "source": [
    "class task1decoder(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "\n",
    "        self.gp = gpsdata()\n",
    "        self.rd = radardata()\n",
    "        self.lc = cameradata()\n",
    "        self.cc = cameradata()\n",
    "        self.rc = cameradata()\n",
    "\n",
    "        if int(num_H/2)*int(num_H/2) > 32:\n",
    "            self.linear = nn.Linear(16*5, int(num_H/2)*int(num_H/2))\n",
    "            self.output_fc = nn.Linear(int(num_H/2)*int(num_H/2), num_H*num_H)\n",
    "            self.output_relu = nn.ReLU()\n",
    "        else:\n",
    "            self.linear = nn.Linear(16*5, 32)\n",
    "            self.output_fc = nn.Linear(32, num_H*num_H)\n",
    "            self.output_relu = nn.ReLU()\n",
    "        \n",
    "        for m in self.modules():\n",
    "            if isinstance(m, (nn.Conv2d, nn.Linear)):\n",
    "                nn.init.xavier_uniform_(m.weight)\n",
    "            elif isinstance(m, nn.BatchNorm2d):\n",
    "                nn.init.constant_(m.weight, 1)\n",
    "                nn.init.constant_(m.bias, 0)\n",
    "    \n",
    "    def forward(self,gps,radar,left_cam,center_cam,right_cam,onoffdict):\n",
    "        bs = gps.size(0)\n",
    "        \n",
    "        if onoffdict['GPS']:\n",
    "             gps_out = self.gp(gps)\n",
    "        else:\n",
    "             gps_out = torch.zeros(bs, 16).to(device)\n",
    "        \n",
    "        if onoffdict['RADAR']:\n",
    "            radar_out = self.rd(radar)\n",
    "        else:\n",
    "            radar_out = torch.zeros(bs, 16).to(device)\n",
    "        \n",
    "        if onoffdict['CAMERAS']:\n",
    "            left_cam, center_cam, right_cam = process_imgs(gps, left_cam, center_cam, right_cam, crop_size = (150,150))\n",
    "            lc_out = self.lc(left_cam)\n",
    "            cc_out = self.cc(center_cam)\n",
    "            rc_out = self.rc(right_cam)\n",
    "        else:\n",
    "            lc_out = torch.zeros(bs, 16).to(device)\n",
    "            cc_out = torch.zeros(bs, 16).to(device)\n",
    "            rc_out = torch.zeros(bs, 16).to(device)\n",
    "\n",
    "        combined = torch.cat((gps_out, radar_out, lc_out, cc_out, rc_out), dim=1)\n",
    "        \n",
    "        output = self.linear(combined)\n",
    "        output = self.output_relu(output)\n",
    "        output = self.output_fc(output)\n",
    "        output = self.output_relu(output)\n",
    "        output = output.view(output.size(0), 1, num_H, num_H)\n",
    "\n",
    "        return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "010e02db",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Decoderwithmsi(nn.Module):\n",
    "    def __init__(self, reduction):\n",
    "        super().__init__()\n",
    "        self.task1decoder = torch.load(task1_weight_path+\"task1Decoder.pth\")\n",
    "        self.gp = self.task1decoder.gp\n",
    "        self.rd = self.task1decoder.rd\n",
    "        self.lc = self.task1decoder.lc\n",
    "        self.cc = self.task1decoder.cc\n",
    "        self.rc = self.task1decoder.rc\n",
    "        self.bde = torch.load(weight_pathb+\"task2Decoder.pth\")\n",
    "        self.allow_update = False  # Initially, do not allow weight updates\n",
    "        #self.bde = baselinedecoder(reduction) \n",
    "\n",
    "        if int(num_H/2)*int(num_H/2) > 32:\n",
    "            self.linear = self.task1decoder.linear\n",
    "            self.output_fc1= nn.Linear(int(num_H/2)*int(num_H/2)+num_H*num_H, 2*num_H*num_H)\n",
    "            self.output_fc2 = nn.Linear(2*num_H*num_H, num_H*num_H)\n",
    "            self.output_relu = nn.ReLU()\n",
    "        else:\n",
    "            self.linear = self.task1decoder.linear\n",
    "            self.output_fc1 = nn.Linear(32+num_H*num_H, 2*num_H*num_H)\n",
    "            self.output_fc2 = nn.Linear(2*num_H*num_H, num_H*num_H)\n",
    "            self.output_relu = nn.ReLU()\n",
    "        \n",
    "        for m in self.modules():\n",
    "            if isinstance(m, (nn.Conv2d, nn.Linear)):\n",
    "                nn.init.xavier_uniform_(m.weight)\n",
    "            elif isinstance(m, nn.BatchNorm2d):\n",
    "                nn.init.constant_(m.weight, 1)\n",
    "                nn.init.constant_(m.bias, 0)\n",
    "    \n",
    "    def forward(self, Hencoded,gps,radar,left_cam,center_cam,right_cam,onoffdict):\n",
    "        bs = Hencoded.size(0)\n",
    "        if self.allow_update:\n",
    "            Hdecoded=self.bde(Hencoded)\n",
    "            if onoffdict['GPS']:\n",
    "                 gps_out = self.gp(gps)\n",
    "            else:\n",
    "                 gps_out = torch.zeros(bs, 16).to(device)\n",
    "\n",
    "            if onoffdict['RADAR']:\n",
    "                radar_out = self.rd(radar)\n",
    "            else:\n",
    "                radar_out = torch.zeros(bs, 16).to(device)\n",
    "\n",
    "            if onoffdict['CAMERAS']:\n",
    "                left_cam, center_cam, right_cam = process_imgs(gps, left_cam, center_cam, right_cam, crop_size = (150,150))\n",
    "                lc_out = self.lc(left_cam)\n",
    "                cc_out = self.cc(center_cam)\n",
    "                rc_out = self.rc(right_cam)\n",
    "            else:\n",
    "                lc_out = torch.zeros(bs, 16).to(device)\n",
    "                cc_out = torch.zeros(bs, 16).to(device)\n",
    "                rc_out = torch.zeros(bs, 16).to(device)\n",
    "\n",
    "            combined1 = torch.cat((gps_out, radar_out, lc_out, cc_out, rc_out), dim=1)\n",
    "\n",
    "            output = self.linear(combined1)\n",
    "        else:\n",
    "            with torch.no_grad():\n",
    "                Hdecoded=self.bde(Hencoded)\n",
    "                if onoffdict['GPS']:\n",
    "                     gps_out = self.gp(gps)\n",
    "                else:\n",
    "                     gps_out = torch.zeros(bs, 16).to(device)\n",
    "\n",
    "                if onoffdict['RADAR']:\n",
    "                    radar_out = self.rd(radar)\n",
    "                else:\n",
    "                    radar_out = torch.zeros(bs, 16).to(device)\n",
    "\n",
    "                if onoffdict['CAMERAS']:\n",
    "                    left_cam, center_cam, right_cam = process_imgs(gps, left_cam, center_cam, right_cam, crop_size = (150,150))\n",
    "                    lc_out = self.lc(left_cam)\n",
    "                    cc_out = self.cc(center_cam)\n",
    "                    rc_out = self.rc(right_cam)\n",
    "                else:\n",
    "                    lc_out = torch.zeros(bs, 16).to(device)\n",
    "                    cc_out = torch.zeros(bs, 16).to(device)\n",
    "                    rc_out = torch.zeros(bs, 16).to(device)\n",
    "\n",
    "                combined1 = torch.cat((gps_out, radar_out, lc_out, cc_out, rc_out), dim=1)\n",
    "\n",
    "                output = self.linear(combined1)\n",
    "                \n",
    "        \n",
    "        \n",
    "        \n",
    "        output = self.output_relu(output)\n",
    "        combined2 = torch.cat((output, Hdecoded.view(bs,-1)), dim=1)\n",
    "        output = self.output_fc1(combined2)\n",
    "        output = self.output_relu(output)\n",
    "        output = self.output_fc2(output)\n",
    "        output = output.view(bs, 1, num_H, num_H)\n",
    "\n",
    "        return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "24be9be7",
   "metadata": {},
   "outputs": [],
   "source": [
    "#complete task 2 model including encoder, decoder and channel\n",
    "class task2model(nn.Module):\n",
    "    def __init__(self, reduction=16):\n",
    "        super().__init__()\n",
    "        \n",
    "        self.en=torch.load(weight_pathb+\"task2Encoder.pth\")\n",
    "        self.allow_update = False  # Initially, do not allow weight updates\n",
    "        \n",
    "        self.de=Decoderwithmsi(reduction)\n",
    "        \n",
    "    \n",
    "   \n",
    "    def forward(self, Hin, gps, radar, left_cam, center_cam, right_cam, onoffdict): \n",
    "        \n",
    "        #Encoder\n",
    "        if self.allow_update:\n",
    "            Hencoded=self.en(Hin)\n",
    "        else:\n",
    "            with torch.no_grad():\n",
    "                Hencoded=self.en(Hin)\n",
    "        \n",
    "        \n",
    "        \n",
    "        \n",
    "        #Decoder\n",
    "        \n",
    "        Hdecoded=self.de(Hencoded,gps,radar,left_cam,center_cam,right_cam,onoffdict)\n",
    "        \n",
    "\n",
    "        return Hdecoded"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b8378ab0",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "f2ba50ec",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Loss\n",
    "\n",
    "#criterion=nn.BCELoss()\n",
    "#criterion = nn.CrossEntropyLoss()\n",
    "criterion= nn.MSELoss().to(device)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0ac6c938",
   "metadata": {},
   "source": [
    "# Inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "09917417",
   "metadata": {},
   "outputs": [],
   "source": [
    "def run_test(model, test_loader, device, criterion):\n",
    "    num_test_batches = len(test_loader)\n",
    "    model.eval()\n",
    "    with torch.no_grad():\n",
    "        mse1 = 0\n",
    "        for b, (X_test, y_test) in enumerate(test_loader):\n",
    "            y_test = y_test.to(device)\n",
    "            Xin = CSI_abs_reshape(X_test[0])\n",
    "            y_pred = model(Xin.to(device),X_test[1].to(device),X_test[2].to(device),X_test[3].to(device),X_test[4].to(device),X_test[5].to(device), onoffdict = onoffdict)\n",
    "            y_test_reshaped = CSI_abs_reshape(y_test)\n",
    "            mse0 = criterion(y_pred, y_test_reshaped) \n",
    "            mse1 += mse0 \n",
    "        \n",
    "    avg_mse = mse1 / num_test_batches\n",
    "    return avg_mse.item()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "id": "910180db",
   "metadata": {},
   "outputs": [],
   "source": [
    "def calculate_confidence_interval(data, confidence=0.95):\n",
    "    n = len(data)\n",
    "    mean = np.mean(data)\n",
    "    se = stats.sem(data)\n",
    "    h = se * stats.t.ppf((1 + confidence) / 2., n-1)\n",
    "    return mean, h"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "f6b9decf",
   "metadata": {},
   "outputs": [],
   "source": [
    "test_dataset = DeepVerseChallengeLoaderTaskTwo(csv_path = r'./dataset_validation.csv')\n",
    "test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=True, num_workers=5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "70a7d594",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "\n",
    "h_list = torch.tensor([])\n",
    "for b, (x,h) in enumerate(test_loader):\n",
    "    h = CSI_abs_reshape(h)\n",
    "    h_list = torch.cat([h_list,h])\n",
    "target_loss = torch.mean((torch.abs(h_list) - torch.mean(torch.abs(h_list))) ** 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "id": "11097b25",
   "metadata": {},
   "outputs": [],
   "source": [
    "num_runs =10"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "id": "1b69466f",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Percentage Improvement Mean Achieved: 56.6364%\n",
      "Percentage Improvement Confidence Interval Achieved: 0.4498%\n",
      "Mean MSE: 0.4838\n",
      "95% Confidence Interval: (0.4788, 0.4889)\n",
      "Margin of Error: 0.0050\n"
     ]
    }
   ],
   "source": [
    "avg_mse_list = []\n",
    "improvement_list = []\n",
    "for _ in range(num_runs):\n",
    "    test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=True, num_workers=5)\n",
    "    model = torch.load(weight_path + \"task2.pth\").to(device)\n",
    "    avg_mse = run_test(model, test_loader, device, criterion)\n",
    "    avg_mse_list.append(avg_mse)\n",
    "    improvement = (target_loss.item() - avg_mse) / target_loss.item() * 100\n",
    "    improvement_list.append(improvement)\n",
    "mean_mse, margin_of_error = calculate_confidence_interval(avg_mse_list)\n",
    "improvement_mean, improvement_margin_of_error = calculate_confidence_interval(improvement_list)\n",
    "print(f'Percentage Improvement Mean Achieved: {improvement_mean:.4f}%')\n",
    "print(f'Percentage Improvement Confidence Interval Achieved: {improvement_margin_of_error:.4f}%')\n",
    "print(f\"Mean MSE: {mean_mse:.4f}\")\n",
    "print(f\"95% Confidence Interval: ({mean_mse - margin_of_error:.4f}, {mean_mse + margin_of_error:.4f})\")\n",
    "print(f\"Margin of Error: {margin_of_error:.4f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "19646553",
   "metadata": {},
   "source": [
    "# Change Dictionary"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "id": "7d02138e",
   "metadata": {},
   "outputs": [],
   "source": [
    "onoffdict={'GPS': True, 'CAMERAS': True, 'RADAR': False}\n",
    "weight_path=f'models/TransNettask2/cr{reduction}/gps{onoffdict[\"GPS\"]}_cam{onoffdict[\"CAMERAS\"]}_rad{onoffdict[\"RADAR\"]}/'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "id": "cd977a09",
   "metadata": {},
   "outputs": [],
   "source": [
    "if not os.path.exists(weight_path):\n",
    "    os.makedirs(weight_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "id": "fc2ca84e",
   "metadata": {},
   "outputs": [],
   "source": [
    "task1_weight_path=f'models/task1/gps{onoffdict[\"GPS\"]}_cam{onoffdict[\"CAMERAS\"]}_rad{onoffdict[\"RADAR\"]}/'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "id": "f43ca2fe",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Decoderwithmsi(nn.Module):\n",
    "    def __init__(self, reduction):\n",
    "        super().__init__()\n",
    "        self.task1decoder = torch.load(task1_weight_path+\"task1Decoder.pth\")\n",
    "        self.gp = self.task1decoder.gp\n",
    "        self.rd = self.task1decoder.rd\n",
    "        self.lc = self.task1decoder.lc\n",
    "        self.cc = self.task1decoder.cc\n",
    "        self.rc = self.task1decoder.rc\n",
    "        self.bde = torch.load(weight_pathb+\"task2Decoder.pth\")\n",
    "        self.allow_update = False  # Initially, do not allow weight updates\n",
    "        #self.bde = baselinedecoder(reduction) \n",
    "\n",
    "        if int(num_H/2)*int(num_H/2) > 32:\n",
    "            self.linear = self.task1decoder.linear\n",
    "            self.output_fc1= nn.Linear(int(num_H/2)*int(num_H/2)+num_H*num_H, 2*num_H*num_H)\n",
    "            self.output_fc2 = nn.Linear(2*num_H*num_H, num_H*num_H)\n",
    "            self.output_relu = nn.ReLU()\n",
    "        else:\n",
    "            self.linear = self.task1decoder.linear\n",
    "            self.output_fc1 = nn.Linear(32+num_H*num_H, 2*num_H*num_H)\n",
    "            self.output_fc2 = nn.Linear(2*num_H*num_H, num_H*num_H)\n",
    "            self.output_relu = nn.ReLU()\n",
    "        \n",
    "        for m in self.modules():\n",
    "            if isinstance(m, (nn.Conv2d, nn.Linear)):\n",
    "                nn.init.xavier_uniform_(m.weight)\n",
    "            elif isinstance(m, nn.BatchNorm2d):\n",
    "                nn.init.constant_(m.weight, 1)\n",
    "                nn.init.constant_(m.bias, 0)\n",
    "    \n",
    "    def forward(self, Hencoded,gps,radar,left_cam,center_cam,right_cam,onoffdict):\n",
    "        bs = Hencoded.size(0)\n",
    "        if self.allow_update:\n",
    "            Hdecoded=self.bde(Hencoded)\n",
    "            if onoffdict['GPS']:\n",
    "                 gps_out = self.gp(gps)\n",
    "            else:\n",
    "                 gps_out = torch.zeros(bs, 16).to(device)\n",
    "\n",
    "            if onoffdict['RADAR']:\n",
    "                radar_out = self.rd(radar)\n",
    "            else:\n",
    "                radar_out = torch.zeros(bs, 16).to(device)\n",
    "\n",
    "            if onoffdict['CAMERAS']:\n",
    "                left_cam, center_cam, right_cam = process_imgs(gps, left_cam, center_cam, right_cam, crop_size = (150,150))\n",
    "                lc_out = self.lc(left_cam)\n",
    "                cc_out = self.cc(center_cam)\n",
    "                rc_out = self.rc(right_cam)\n",
    "            else:\n",
    "                lc_out = torch.zeros(bs, 16).to(device)\n",
    "                cc_out = torch.zeros(bs, 16).to(device)\n",
    "                rc_out = torch.zeros(bs, 16).to(device)\n",
    "\n",
    "            combined1 = torch.cat((gps_out, radar_out, lc_out, cc_out, rc_out), dim=1)\n",
    "\n",
    "            output = self.linear(combined1)\n",
    "        else:\n",
    "            with torch.no_grad():\n",
    "                Hdecoded=self.bde(Hencoded)\n",
    "                if onoffdict['GPS']:\n",
    "                     gps_out = self.gp(gps)\n",
    "                else:\n",
    "                     gps_out = torch.zeros(bs, 16).to(device)\n",
    "\n",
    "                if onoffdict['RADAR']:\n",
    "                    radar_out = self.rd(radar)\n",
    "                else:\n",
    "                    radar_out = torch.zeros(bs, 16).to(device)\n",
    "\n",
    "                if onoffdict['CAMERAS']:\n",
    "                    left_cam, center_cam, right_cam = process_imgs(gps, left_cam, center_cam, right_cam, crop_size = (150,150))\n",
    "                    lc_out = self.lc(left_cam)\n",
    "                    cc_out = self.cc(center_cam)\n",
    "                    rc_out = self.rc(right_cam)\n",
    "                else:\n",
    "                    lc_out = torch.zeros(bs, 16).to(device)\n",
    "                    cc_out = torch.zeros(bs, 16).to(device)\n",
    "                    rc_out = torch.zeros(bs, 16).to(device)\n",
    "\n",
    "                combined1 = torch.cat((gps_out, radar_out, lc_out, cc_out, rc_out), dim=1)\n",
    "\n",
    "                output = self.linear(combined1)\n",
    "                \n",
    "        \n",
    "        \n",
    "        \n",
    "        output = self.output_relu(output)\n",
    "        combined2 = torch.cat((output, Hdecoded.view(bs,-1)), dim=1)\n",
    "        output = self.output_fc1(combined2)\n",
    "        output = self.output_relu(output)\n",
    "        output = self.output_fc2(output)\n",
    "        output = output.view(bs, 1, num_H, num_H)\n",
    "\n",
    "        return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "id": "4fe56827",
   "metadata": {},
   "outputs": [],
   "source": [
    "#complete task 2 model including encoder, decoder and channel\n",
    "class task2model(nn.Module):\n",
    "    def __init__(self, reduction=16):\n",
    "        super().__init__()\n",
    "        \n",
    "        self.en=torch.load(weight_pathb+\"task2Encoder.pth\")\n",
    "        self.allow_update = False  # Initially, do not allow weight updates\n",
    "        \n",
    "        self.de=Decoderwithmsi(reduction)\n",
    "        \n",
    "    \n",
    "   \n",
    "    def forward(self, Hin, gps, radar, left_cam, center_cam, right_cam, onoffdict): \n",
    "        \n",
    "        #Encoder\n",
    "        if self.allow_update:\n",
    "            Hencoded=self.en(Hin)\n",
    "        else:\n",
    "            with torch.no_grad():\n",
    "                Hencoded=self.en(Hin)\n",
    "        \n",
    "        \n",
    "        \n",
    "        \n",
    "        #Decoder\n",
    "        \n",
    "        Hdecoded=self.de(Hencoded,gps,radar,left_cam,center_cam,right_cam,onoffdict)\n",
    "        \n",
    "\n",
    "        return Hdecoded"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9dd0d773",
   "metadata": {},
   "source": [
    "# Inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "id": "b4db6744",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Percentage Improvement Mean Achieved: 56.0821%\n",
      "Percentage Improvement Confidence Interval Achieved: 0.6571%\n",
      "Mean MSE: 0.4900\n",
      "95% Confidence Interval: (0.4827, 0.4974)\n",
      "Margin of Error: 0.0073\n"
     ]
    }
   ],
   "source": [
    "avg_mse_list = []\n",
    "improvement_list = []\n",
    "for _ in range(num_runs):\n",
    "    test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=True, num_workers=5)\n",
    "    model = torch.load(weight_path + \"task2.pth\").to(device)\n",
    "    avg_mse = run_test(model, test_loader, device, criterion)\n",
    "    avg_mse_list.append(avg_mse)\n",
    "    improvement = (target_loss.item() - avg_mse) / target_loss.item() * 100\n",
    "    improvement_list.append(improvement)\n",
    "mean_mse, margin_of_error = calculate_confidence_interval(avg_mse_list)\n",
    "improvement_mean, improvement_margin_of_error = calculate_confidence_interval(improvement_list)\n",
    "print(f'Percentage Improvement Mean Achieved: {improvement_mean:.4f}%')\n",
    "print(f'Percentage Improvement Confidence Interval Achieved: {improvement_margin_of_error:.4f}%')\n",
    "print(f\"Mean MSE: {mean_mse:.4f}\")\n",
    "print(f\"95% Confidence Interval: ({mean_mse - margin_of_error:.4f}, {mean_mse + margin_of_error:.4f})\")\n",
    "print(f\"Margin of Error: {margin_of_error:.4f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6f43022b",
   "metadata": {},
   "source": [
    "# Change Dictionary"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "id": "b0908d8e",
   "metadata": {},
   "outputs": [],
   "source": [
    "onoffdict={'GPS': True, 'CAMERAS': False, 'RADAR': True}\n",
    "weight_path=f'models/TransNettask2/cr{reduction}/gps{onoffdict[\"GPS\"]}_cam{onoffdict[\"CAMERAS\"]}_rad{onoffdict[\"RADAR\"]}/'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "id": "0c1182d1",
   "metadata": {},
   "outputs": [],
   "source": [
    "if not os.path.exists(weight_path):\n",
    "    os.makedirs(weight_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "id": "d206c2f8",
   "metadata": {},
   "outputs": [],
   "source": [
    "task1_weight_path=f'models/task1/gps{onoffdict[\"GPS\"]}_cam{onoffdict[\"CAMERAS\"]}_rad{onoffdict[\"RADAR\"]}/'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "id": "6e2782ce",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Decoderwithmsi(nn.Module):\n",
    "    def __init__(self, reduction):\n",
    "        super().__init__()\n",
    "        self.task1decoder = torch.load(task1_weight_path+\"task1Decoder.pth\")\n",
    "        self.gp = self.task1decoder.gp\n",
    "        self.rd = self.task1decoder.rd\n",
    "        self.lc = self.task1decoder.lc\n",
    "        self.cc = self.task1decoder.cc\n",
    "        self.rc = self.task1decoder.rc\n",
    "        self.bde = torch.load(weight_pathb+\"task2Decoder.pth\")\n",
    "        self.allow_update = False  # Initially, do not allow weight updates\n",
    "        #self.bde = baselinedecoder(reduction) \n",
    "\n",
    "        if int(num_H/2)*int(num_H/2) > 32:\n",
    "            self.linear = self.task1decoder.linear\n",
    "            self.output_fc1= nn.Linear(int(num_H/2)*int(num_H/2)+num_H*num_H, 2*num_H*num_H)\n",
    "            self.output_fc2 = nn.Linear(2*num_H*num_H, num_H*num_H)\n",
    "            self.output_relu = nn.ReLU()\n",
    "        else:\n",
    "            self.linear = self.task1decoder.linear\n",
    "            self.output_fc1 = nn.Linear(32+num_H*num_H, 2*num_H*num_H)\n",
    "            self.output_fc2 = nn.Linear(2*num_H*num_H, num_H*num_H)\n",
    "            self.output_relu = nn.ReLU()\n",
    "        \n",
    "        for m in self.modules():\n",
    "            if isinstance(m, (nn.Conv2d, nn.Linear)):\n",
    "                nn.init.xavier_uniform_(m.weight)\n",
    "            elif isinstance(m, nn.BatchNorm2d):\n",
    "                nn.init.constant_(m.weight, 1)\n",
    "                nn.init.constant_(m.bias, 0)\n",
    "    \n",
    "    def forward(self, Hencoded,gps,radar,left_cam,center_cam,right_cam,onoffdict):\n",
    "        bs = Hencoded.size(0)\n",
    "        if self.allow_update:\n",
    "            Hdecoded=self.bde(Hencoded)\n",
    "            if onoffdict['GPS']:\n",
    "                 gps_out = self.gp(gps)\n",
    "            else:\n",
    "                 gps_out = torch.zeros(bs, 16).to(device)\n",
    "\n",
    "            if onoffdict['RADAR']:\n",
    "                radar_out = self.rd(radar)\n",
    "            else:\n",
    "                radar_out = torch.zeros(bs, 16).to(device)\n",
    "\n",
    "            if onoffdict['CAMERAS']:\n",
    "                left_cam, center_cam, right_cam = process_imgs(gps, left_cam, center_cam, right_cam, crop_size = (150,150))\n",
    "                lc_out = self.lc(left_cam)\n",
    "                cc_out = self.cc(center_cam)\n",
    "                rc_out = self.rc(right_cam)\n",
    "            else:\n",
    "                lc_out = torch.zeros(bs, 16).to(device)\n",
    "                cc_out = torch.zeros(bs, 16).to(device)\n",
    "                rc_out = torch.zeros(bs, 16).to(device)\n",
    "\n",
    "            combined1 = torch.cat((gps_out, radar_out, lc_out, cc_out, rc_out), dim=1)\n",
    "\n",
    "            output = self.linear(combined1)\n",
    "        else:\n",
    "            with torch.no_grad():\n",
    "                Hdecoded=self.bde(Hencoded)\n",
    "                if onoffdict['GPS']:\n",
    "                     gps_out = self.gp(gps)\n",
    "                else:\n",
    "                     gps_out = torch.zeros(bs, 16).to(device)\n",
    "\n",
    "                if onoffdict['RADAR']:\n",
    "                    radar_out = self.rd(radar)\n",
    "                else:\n",
    "                    radar_out = torch.zeros(bs, 16).to(device)\n",
    "\n",
    "                if onoffdict['CAMERAS']:\n",
    "                    left_cam, center_cam, right_cam = process_imgs(gps, left_cam, center_cam, right_cam, crop_size = (150,150))\n",
    "                    lc_out = self.lc(left_cam)\n",
    "                    cc_out = self.cc(center_cam)\n",
    "                    rc_out = self.rc(right_cam)\n",
    "                else:\n",
    "                    lc_out = torch.zeros(bs, 16).to(device)\n",
    "                    cc_out = torch.zeros(bs, 16).to(device)\n",
    "                    rc_out = torch.zeros(bs, 16).to(device)\n",
    "\n",
    "                combined1 = torch.cat((gps_out, radar_out, lc_out, cc_out, rc_out), dim=1)\n",
    "\n",
    "                output = self.linear(combined1)\n",
    "                \n",
    "        \n",
    "        \n",
    "        \n",
    "        output = self.output_relu(output)\n",
    "        combined2 = torch.cat((output, Hdecoded.view(bs,-1)), dim=1)\n",
    "        output = self.output_fc1(combined2)\n",
    "        output = self.output_relu(output)\n",
    "        output = self.output_fc2(output)\n",
    "        output = output.view(bs, 1, num_H, num_H)\n",
    "\n",
    "        return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "id": "b6aa587b",
   "metadata": {},
   "outputs": [],
   "source": [
    "#complete task 2 model including encoder, decoder and channel\n",
    "class task2model(nn.Module):\n",
    "    def __init__(self, reduction=16):\n",
    "        super().__init__()\n",
    "        \n",
    "        self.en=torch.load(weight_pathb+\"task2Encoder.pth\")\n",
    "        self.allow_update = False  # Initially, do not allow weight updates\n",
    "        \n",
    "        self.de=Decoderwithmsi(reduction)\n",
    "        \n",
    "    \n",
    "   \n",
    "    def forward(self, Hin, gps, radar, left_cam, center_cam, right_cam, onoffdict): \n",
    "        \n",
    "        #Encoder\n",
    "        if self.allow_update:\n",
    "            Hencoded=self.en(Hin)\n",
    "        else:\n",
    "            with torch.no_grad():\n",
    "                Hencoded=self.en(Hin)\n",
    "        \n",
    "        \n",
    "        \n",
    "        \n",
    "        #Decoder\n",
    "        \n",
    "        Hdecoded=self.de(Hencoded,gps,radar,left_cam,center_cam,right_cam,onoffdict)\n",
    "        \n",
    "\n",
    "        return Hdecoded"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "id": "ef64e987",
   "metadata": {},
   "outputs": [],
   "source": [
    "model=task2model(reduction)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "55d649d7",
   "metadata": {},
   "source": [
    "# Inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "id": "3cc775e3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Percentage Improvement Mean Achieved: 56.0597%\n",
      "Percentage Improvement Confidence Interval Achieved: 0.4678%\n",
      "Mean MSE: 0.4903\n",
      "95% Confidence Interval: (0.4851, 0.4955)\n",
      "Margin of Error: 0.0052\n"
     ]
    }
   ],
   "source": [
    "avg_mse_list = []\n",
    "improvement_list = []\n",
    "for _ in range(num_runs):\n",
    "    test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=True, num_workers=5)\n",
    "    model = torch.load(weight_path + \"task2.pth\").to(device)\n",
    "    avg_mse = run_test(model, test_loader, device, criterion)\n",
    "    avg_mse_list.append(avg_mse)\n",
    "    improvement = (target_loss.item() - avg_mse) / target_loss.item() * 100\n",
    "    improvement_list.append(improvement)\n",
    "mean_mse, margin_of_error = calculate_confidence_interval(avg_mse_list)\n",
    "improvement_mean, improvement_margin_of_error = calculate_confidence_interval(improvement_list)\n",
    "print(f'Percentage Improvement Mean Achieved: {improvement_mean:.4f}%')\n",
    "print(f'Percentage Improvement Confidence Interval Achieved: {improvement_margin_of_error:.4f}%')\n",
    "print(f\"Mean MSE: {mean_mse:.4f}\")\n",
    "print(f\"95% Confidence Interval: ({mean_mse - margin_of_error:.4f}, {mean_mse + margin_of_error:.4f})\")\n",
    "print(f\"Margin of Error: {margin_of_error:.4f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2fae1529",
   "metadata": {},
   "source": [
    "# Change Dictionary"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "id": "f2f72743",
   "metadata": {},
   "outputs": [],
   "source": [
    "onoffdict={'GPS': True, 'CAMERAS': False, 'RADAR': False}\n",
    "weight_path=f'models/TransNettask2/cr{reduction}/gps{onoffdict[\"GPS\"]}_cam{onoffdict[\"CAMERAS\"]}_rad{onoffdict[\"RADAR\"]}/'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "id": "02e87ce4",
   "metadata": {},
   "outputs": [],
   "source": [
    "if not os.path.exists(weight_path):\n",
    "    os.makedirs(weight_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "id": "bb8b49c1",
   "metadata": {},
   "outputs": [],
   "source": [
    "task1_weight_path=f'models/task1/gps{onoffdict[\"GPS\"]}_cam{onoffdict[\"CAMERAS\"]}_rad{onoffdict[\"RADAR\"]}/'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "id": "ce1cefaf",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Decoderwithmsi(nn.Module):\n",
    "    def __init__(self, reduction):\n",
    "        super().__init__()\n",
    "        self.task1decoder = torch.load(task1_weight_path+\"task1Decoder.pth\")\n",
    "        self.gp = self.task1decoder.gp\n",
    "        self.rd = self.task1decoder.rd\n",
    "        self.lc = self.task1decoder.lc\n",
    "        self.cc = self.task1decoder.cc\n",
    "        self.rc = self.task1decoder.rc\n",
    "        self.bde = torch.load(weight_pathb+\"task2Decoder.pth\")\n",
    "        self.allow_update = False  # Initially, do not allow weight updates\n",
    "        #self.bde = baselinedecoder(reduction) \n",
    "\n",
    "        if int(num_H/2)*int(num_H/2) > 32:\n",
    "            self.linear = self.task1decoder.linear\n",
    "            self.output_fc1= nn.Linear(int(num_H/2)*int(num_H/2)+num_H*num_H, 2*num_H*num_H)\n",
    "            self.output_fc2 = nn.Linear(2*num_H*num_H, num_H*num_H)\n",
    "            self.output_relu = nn.ReLU()\n",
    "        else:\n",
    "            self.linear = self.task1decoder.linear\n",
    "            self.output_fc1 = nn.Linear(32+num_H*num_H, 2*num_H*num_H)\n",
    "            self.output_fc2 = nn.Linear(2*num_H*num_H, num_H*num_H)\n",
    "            self.output_relu = nn.ReLU()\n",
    "        \n",
    "        for m in self.modules():\n",
    "            if isinstance(m, (nn.Conv2d, nn.Linear)):\n",
    "                nn.init.xavier_uniform_(m.weight)\n",
    "            elif isinstance(m, nn.BatchNorm2d):\n",
    "                nn.init.constant_(m.weight, 1)\n",
    "                nn.init.constant_(m.bias, 0)\n",
    "    \n",
    "    def forward(self, Hencoded,gps,radar,left_cam,center_cam,right_cam,onoffdict):\n",
    "        bs = Hencoded.size(0)\n",
    "        if self.allow_update:\n",
    "            Hdecoded=self.bde(Hencoded)\n",
    "            if onoffdict['GPS']:\n",
    "                 gps_out = self.gp(gps)\n",
    "            else:\n",
    "                 gps_out = torch.zeros(bs, 16).to(device)\n",
    "\n",
    "            if onoffdict['RADAR']:\n",
    "                radar_out = self.rd(radar)\n",
    "            else:\n",
    "                radar_out = torch.zeros(bs, 16).to(device)\n",
    "\n",
    "            if onoffdict['CAMERAS']:\n",
    "                left_cam, center_cam, right_cam = process_imgs(gps, left_cam, center_cam, right_cam, crop_size = (150,150))\n",
    "                lc_out = self.lc(left_cam)\n",
    "                cc_out = self.cc(center_cam)\n",
    "                rc_out = self.rc(right_cam)\n",
    "            else:\n",
    "                lc_out = torch.zeros(bs, 16).to(device)\n",
    "                cc_out = torch.zeros(bs, 16).to(device)\n",
    "                rc_out = torch.zeros(bs, 16).to(device)\n",
    "\n",
    "            combined1 = torch.cat((gps_out, radar_out, lc_out, cc_out, rc_out), dim=1)\n",
    "\n",
    "            output = self.linear(combined1)\n",
    "        else:\n",
    "            with torch.no_grad():\n",
    "                Hdecoded=self.bde(Hencoded)\n",
    "                if onoffdict['GPS']:\n",
    "                     gps_out = self.gp(gps)\n",
    "                else:\n",
    "                     gps_out = torch.zeros(bs, 16).to(device)\n",
    "\n",
    "                if onoffdict['RADAR']:\n",
    "                    radar_out = self.rd(radar)\n",
    "                else:\n",
    "                    radar_out = torch.zeros(bs, 16).to(device)\n",
    "\n",
    "                if onoffdict['CAMERAS']:\n",
    "                    left_cam, center_cam, right_cam = process_imgs(gps, left_cam, center_cam, right_cam, crop_size = (150,150))\n",
    "                    lc_out = self.lc(left_cam)\n",
    "                    cc_out = self.cc(center_cam)\n",
    "                    rc_out = self.rc(right_cam)\n",
    "                else:\n",
    "                    lc_out = torch.zeros(bs, 16).to(device)\n",
    "                    cc_out = torch.zeros(bs, 16).to(device)\n",
    "                    rc_out = torch.zeros(bs, 16).to(device)\n",
    "\n",
    "                combined1 = torch.cat((gps_out, radar_out, lc_out, cc_out, rc_out), dim=1)\n",
    "\n",
    "                output = self.linear(combined1)\n",
    "                \n",
    "        \n",
    "        \n",
    "        \n",
    "        output = self.output_relu(output)\n",
    "        combined2 = torch.cat((output, Hdecoded.view(bs,-1)), dim=1)\n",
    "        output = self.output_fc1(combined2)\n",
    "        output = self.output_relu(output)\n",
    "        output = self.output_fc2(output)\n",
    "        output = output.view(bs, 1, num_H, num_H)\n",
    "\n",
    "        return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "id": "7e685068",
   "metadata": {},
   "outputs": [],
   "source": [
    "#complete task 2 model including encoder, decoder and channel\n",
    "class task2model(nn.Module):\n",
    "    def __init__(self, reduction=16):\n",
    "        super().__init__()\n",
    "        \n",
    "        self.en=torch.load(weight_pathb+\"task2Encoder.pth\")\n",
    "        self.allow_update = False  # Initially, do not allow weight updates\n",
    "        \n",
    "        self.de=Decoderwithmsi(reduction)\n",
    "        \n",
    "    \n",
    "   \n",
    "    def forward(self, Hin, gps, radar, left_cam, center_cam, right_cam, onoffdict): \n",
    "        \n",
    "        #Encoder\n",
    "        if self.allow_update:\n",
    "            Hencoded=self.en(Hin)\n",
    "        else:\n",
    "            with torch.no_grad():\n",
    "                Hencoded=self.en(Hin)\n",
    "        \n",
    "        \n",
    "        \n",
    "        \n",
    "        #Decoder\n",
    "        \n",
    "        Hdecoded=self.de(Hencoded,gps,radar,left_cam,center_cam,right_cam,onoffdict)\n",
    "        \n",
    "\n",
    "        return Hdecoded"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b6efa9a6",
   "metadata": {},
   "source": [
    "# Inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "id": "f5c88883",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Percentage Improvement Mean Achieved: 55.0498%\n",
      "Percentage Improvement Confidence Interval Achieved: 0.4695%\n",
      "Mean MSE: 0.5015\n",
      "95% Confidence Interval: (0.4963, 0.5068)\n",
      "Margin of Error: 0.0052\n"
     ]
    }
   ],
   "source": [
    "avg_mse_list = []\n",
    "improvement_list = []\n",
    "for _ in range(num_runs):\n",
    "    test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=True, num_workers=5)\n",
    "    model = torch.load(weight_path + \"task2.pth\").to(device)\n",
    "    avg_mse = run_test(model, test_loader, device, criterion)\n",
    "    avg_mse_list.append(avg_mse)\n",
    "    improvement = (target_loss.item() - avg_mse) / target_loss.item() * 100\n",
    "    improvement_list.append(improvement)\n",
    "mean_mse, margin_of_error = calculate_confidence_interval(avg_mse_list)\n",
    "improvement_mean, improvement_margin_of_error = calculate_confidence_interval(improvement_list)\n",
    "print(f'Percentage Improvement Mean Achieved: {improvement_mean:.4f}%')\n",
    "print(f'Percentage Improvement Confidence Interval Achieved: {improvement_margin_of_error:.4f}%')\n",
    "print(f\"Mean MSE: {mean_mse:.4f}\")\n",
    "print(f\"95% Confidence Interval: ({mean_mse - margin_of_error:.4f}, {mean_mse + margin_of_error:.4f})\")\n",
    "print(f\"Margin of Error: {margin_of_error:.4f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "64b91f7e",
   "metadata": {},
   "source": [
    "# Change Dictionary"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "id": "4d81898f",
   "metadata": {},
   "outputs": [],
   "source": [
    "onoffdict={'GPS': False, 'CAMERAS': True, 'RADAR': True}\n",
    "weight_path=f'models/TransNettask2/cr{reduction}/gps{onoffdict[\"GPS\"]}_cam{onoffdict[\"CAMERAS\"]}_rad{onoffdict[\"RADAR\"]}/'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "id": "4b66eee9",
   "metadata": {},
   "outputs": [],
   "source": [
    "if not os.path.exists(weight_path):\n",
    "    os.makedirs(weight_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "id": "c77e6892",
   "metadata": {},
   "outputs": [],
   "source": [
    "task1_weight_path=f'models/task1/gps{onoffdict[\"GPS\"]}_cam{onoffdict[\"CAMERAS\"]}_rad{onoffdict[\"RADAR\"]}/'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "id": "4d0efe20",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Decoderwithmsi(nn.Module):\n",
    "    def __init__(self, reduction):\n",
    "        super().__init__()\n",
    "        self.task1decoder = torch.load(task1_weight_path+\"task1Decoder.pth\")\n",
    "        self.gp = self.task1decoder.gp\n",
    "        self.rd = self.task1decoder.rd\n",
    "        self.lc = self.task1decoder.lc\n",
    "        self.cc = self.task1decoder.cc\n",
    "        self.rc = self.task1decoder.rc\n",
    "        self.bde = torch.load(weight_pathb+\"task2Decoder.pth\")\n",
    "        self.allow_update = False  # Initially, do not allow weight updates\n",
    "        #self.bde = baselinedecoder(reduction) \n",
    "\n",
    "        if int(num_H/2)*int(num_H/2) > 32:\n",
    "            self.linear = self.task1decoder.linear\n",
    "            self.output_fc1= nn.Linear(int(num_H/2)*int(num_H/2)+num_H*num_H, 2*num_H*num_H)\n",
    "            self.output_fc2 = nn.Linear(2*num_H*num_H, num_H*num_H)\n",
    "            self.output_relu = nn.ReLU()\n",
    "        else:\n",
    "            self.linear = self.task1decoder.linear\n",
    "            self.output_fc1 = nn.Linear(32+num_H*num_H, 2*num_H*num_H)\n",
    "            self.output_fc2 = nn.Linear(2*num_H*num_H, num_H*num_H)\n",
    "            self.output_relu = nn.ReLU()\n",
    "        \n",
    "        for m in self.modules():\n",
    "            if isinstance(m, (nn.Conv2d, nn.Linear)):\n",
    "                nn.init.xavier_uniform_(m.weight)\n",
    "            elif isinstance(m, nn.BatchNorm2d):\n",
    "                nn.init.constant_(m.weight, 1)\n",
    "                nn.init.constant_(m.bias, 0)\n",
    "    \n",
    "    def forward(self, Hencoded,gps,radar,left_cam,center_cam,right_cam,onoffdict):\n",
    "        bs = Hencoded.size(0)\n",
    "        if self.allow_update:\n",
    "            Hdecoded=self.bde(Hencoded)\n",
    "            if onoffdict['GPS']:\n",
    "                 gps_out = self.gp(gps)\n",
    "            else:\n",
    "                 gps_out = torch.zeros(bs, 16).to(device)\n",
    "\n",
    "            if onoffdict['RADAR']:\n",
    "                radar_out = self.rd(radar)\n",
    "            else:\n",
    "                radar_out = torch.zeros(bs, 16).to(device)\n",
    "\n",
    "            if onoffdict['CAMERAS']:\n",
    "                left_cam, center_cam, right_cam = process_imgs(gps, left_cam, center_cam, right_cam, crop_size = (150,150))\n",
    "                lc_out = self.lc(left_cam)\n",
    "                cc_out = self.cc(center_cam)\n",
    "                rc_out = self.rc(right_cam)\n",
    "            else:\n",
    "                lc_out = torch.zeros(bs, 16).to(device)\n",
    "                cc_out = torch.zeros(bs, 16).to(device)\n",
    "                rc_out = torch.zeros(bs, 16).to(device)\n",
    "\n",
    "            combined1 = torch.cat((gps_out, radar_out, lc_out, cc_out, rc_out), dim=1)\n",
    "\n",
    "            output = self.linear(combined1)\n",
    "        else:\n",
    "            with torch.no_grad():\n",
    "                Hdecoded=self.bde(Hencoded)\n",
    "                if onoffdict['GPS']:\n",
    "                     gps_out = self.gp(gps)\n",
    "                else:\n",
    "                     gps_out = torch.zeros(bs, 16).to(device)\n",
    "\n",
    "                if onoffdict['RADAR']:\n",
    "                    radar_out = self.rd(radar)\n",
    "                else:\n",
    "                    radar_out = torch.zeros(bs, 16).to(device)\n",
    "\n",
    "                if onoffdict['CAMERAS']:\n",
    "                    left_cam, center_cam, right_cam = process_imgs(gps, left_cam, center_cam, right_cam, crop_size = (150,150))\n",
    "                    lc_out = self.lc(left_cam)\n",
    "                    cc_out = self.cc(center_cam)\n",
    "                    rc_out = self.rc(right_cam)\n",
    "                else:\n",
    "                    lc_out = torch.zeros(bs, 16).to(device)\n",
    "                    cc_out = torch.zeros(bs, 16).to(device)\n",
    "                    rc_out = torch.zeros(bs, 16).to(device)\n",
    "\n",
    "                combined1 = torch.cat((gps_out, radar_out, lc_out, cc_out, rc_out), dim=1)\n",
    "\n",
    "                output = self.linear(combined1)\n",
    "                \n",
    "        \n",
    "        \n",
    "        \n",
    "        output = self.output_relu(output)\n",
    "        combined2 = torch.cat((output, Hdecoded.view(bs,-1)), dim=1)\n",
    "        output = self.output_fc1(combined2)\n",
    "        output = self.output_relu(output)\n",
    "        output = self.output_fc2(output)\n",
    "        output = output.view(bs, 1, num_H, num_H)\n",
    "\n",
    "        return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "id": "bf923e7d",
   "metadata": {},
   "outputs": [],
   "source": [
    "#complete task 2 model including encoder, decoder and channel\n",
    "class task2model(nn.Module):\n",
    "    def __init__(self, reduction=16):\n",
    "        super().__init__()\n",
    "        \n",
    "        self.en=torch.load(weight_pathb+\"task2Encoder.pth\")\n",
    "        self.allow_update = False  # Initially, do not allow weight updates\n",
    "        \n",
    "        self.de=Decoderwithmsi(reduction)\n",
    "        \n",
    "    \n",
    "   \n",
    "    def forward(self, Hin, gps, radar, left_cam, center_cam, right_cam, onoffdict): \n",
    "        \n",
    "        #Encoder\n",
    "        if self.allow_update:\n",
    "            Hencoded=self.en(Hin)\n",
    "        else:\n",
    "            with torch.no_grad():\n",
    "                Hencoded=self.en(Hin)\n",
    "        \n",
    "        \n",
    "        \n",
    "        \n",
    "        #Decoder\n",
    "        \n",
    "        Hdecoded=self.de(Hencoded,gps,radar,left_cam,center_cam,right_cam,onoffdict)\n",
    "        \n",
    "\n",
    "        return Hdecoded"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "id": "8eed0aa1",
   "metadata": {},
   "outputs": [],
   "source": [
    "model=task2model(reduction)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "id": "618d457b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 70,
   "id": "98f2500b",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Loss\n",
    "\n",
    "#criterion=nn.BCELoss()\n",
    "#criterion = nn.CrossEntropyLoss()\n",
    "criterion= nn.MSELoss().to(device)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5b824d73",
   "metadata": {},
   "source": [
    "# Inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 71,
   "id": "e3c2a393",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Percentage Improvement Mean Achieved: 53.5477%\n",
      "Percentage Improvement Confidence Interval Achieved: 0.6450%\n",
      "Mean MSE: 0.5183\n",
      "95% Confidence Interval: (0.5111, 0.5255)\n",
      "Margin of Error: 0.0072\n"
     ]
    }
   ],
   "source": [
    "avg_mse_list = []\n",
    "improvement_list = []\n",
    "for _ in range(num_runs):\n",
    "    test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=True, num_workers=5)\n",
    "    model = torch.load(weight_path + \"task2.pth\").to(device)\n",
    "    avg_mse = run_test(model, test_loader, device, criterion)\n",
    "    avg_mse_list.append(avg_mse)\n",
    "    improvement = (target_loss.item() - avg_mse) / target_loss.item() * 100\n",
    "    improvement_list.append(improvement)\n",
    "mean_mse, margin_of_error = calculate_confidence_interval(avg_mse_list)\n",
    "improvement_mean, improvement_margin_of_error = calculate_confidence_interval(improvement_list)\n",
    "print(f'Percentage Improvement Mean Achieved: {improvement_mean:.4f}%')\n",
    "print(f'Percentage Improvement Confidence Interval Achieved: {improvement_margin_of_error:.4f}%')\n",
    "print(f\"Mean MSE: {mean_mse:.4f}\")\n",
    "print(f\"95% Confidence Interval: ({mean_mse - margin_of_error:.4f}, {mean_mse + margin_of_error:.4f})\")\n",
    "print(f\"Margin of Error: {margin_of_error:.4f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "20c3297f",
   "metadata": {},
   "source": [
    "# Change Dictionary"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 72,
   "id": "13948651",
   "metadata": {},
   "outputs": [],
   "source": [
    "onoffdict={'GPS': False, 'CAMERAS': True, 'RADAR': False}\n",
    "weight_path=f'models/TransNettask2/cr{reduction}/gps{onoffdict[\"GPS\"]}_cam{onoffdict[\"CAMERAS\"]}_rad{onoffdict[\"RADAR\"]}/'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 73,
   "id": "1d041e63",
   "metadata": {},
   "outputs": [],
   "source": [
    "if not os.path.exists(weight_path):\n",
    "    os.makedirs(weight_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 74,
   "id": "d44e3c7e",
   "metadata": {},
   "outputs": [],
   "source": [
    "task1_weight_path=f'models/task1/gps{onoffdict[\"GPS\"]}_cam{onoffdict[\"CAMERAS\"]}_rad{onoffdict[\"RADAR\"]}/'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 75,
   "id": "03d2bd13",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Decoderwithmsi(nn.Module):\n",
    "    def __init__(self, reduction):\n",
    "        super().__init__()\n",
    "        self.task1decoder = torch.load(task1_weight_path+\"task1Decoder.pth\")\n",
    "        self.gp = self.task1decoder.gp\n",
    "        self.rd = self.task1decoder.rd\n",
    "        self.lc = self.task1decoder.lc\n",
    "        self.cc = self.task1decoder.cc\n",
    "        self.rc = self.task1decoder.rc\n",
    "        self.bde = torch.load(weight_pathb+\"task2Decoder.pth\")\n",
    "        self.allow_update = False  # Initially, do not allow weight updates\n",
    "        #self.bde = baselinedecoder(reduction) \n",
    "\n",
    "        if int(num_H/2)*int(num_H/2) > 32:\n",
    "            self.linear = self.task1decoder.linear\n",
    "            self.output_fc1= nn.Linear(int(num_H/2)*int(num_H/2)+num_H*num_H, 2*num_H*num_H)\n",
    "            self.output_fc2 = nn.Linear(2*num_H*num_H, num_H*num_H)\n",
    "            self.output_relu = nn.ReLU()\n",
    "        else:\n",
    "            self.linear = self.task1decoder.linear\n",
    "            self.output_fc1 = nn.Linear(32+num_H*num_H, 2*num_H*num_H)\n",
    "            self.output_fc2 = nn.Linear(2*num_H*num_H, num_H*num_H)\n",
    "            self.output_relu = nn.ReLU()\n",
    "        \n",
    "        for m in self.modules():\n",
    "            if isinstance(m, (nn.Conv2d, nn.Linear)):\n",
    "                nn.init.xavier_uniform_(m.weight)\n",
    "            elif isinstance(m, nn.BatchNorm2d):\n",
    "                nn.init.constant_(m.weight, 1)\n",
    "                nn.init.constant_(m.bias, 0)\n",
    "    \n",
    "    def forward(self, Hencoded,gps,radar,left_cam,center_cam,right_cam,onoffdict):\n",
    "        bs = Hencoded.size(0)\n",
    "        if self.allow_update:\n",
    "            Hdecoded=self.bde(Hencoded)\n",
    "            if onoffdict['GPS']:\n",
    "                 gps_out = self.gp(gps)\n",
    "            else:\n",
    "                 gps_out = torch.zeros(bs, 16).to(device)\n",
    "\n",
    "            if onoffdict['RADAR']:\n",
    "                radar_out = self.rd(radar)\n",
    "            else:\n",
    "                radar_out = torch.zeros(bs, 16).to(device)\n",
    "\n",
    "            if onoffdict['CAMERAS']:\n",
    "                left_cam, center_cam, right_cam = process_imgs(gps, left_cam, center_cam, right_cam, crop_size = (150,150))\n",
    "                lc_out = self.lc(left_cam)\n",
    "                cc_out = self.cc(center_cam)\n",
    "                rc_out = self.rc(right_cam)\n",
    "            else:\n",
    "                lc_out = torch.zeros(bs, 16).to(device)\n",
    "                cc_out = torch.zeros(bs, 16).to(device)\n",
    "                rc_out = torch.zeros(bs, 16).to(device)\n",
    "\n",
    "            combined1 = torch.cat((gps_out, radar_out, lc_out, cc_out, rc_out), dim=1)\n",
    "\n",
    "            output = self.linear(combined1)\n",
    "        else:\n",
    "            with torch.no_grad():\n",
    "                Hdecoded=self.bde(Hencoded)\n",
    "                if onoffdict['GPS']:\n",
    "                     gps_out = self.gp(gps)\n",
    "                else:\n",
    "                     gps_out = torch.zeros(bs, 16).to(device)\n",
    "\n",
    "                if onoffdict['RADAR']:\n",
    "                    radar_out = self.rd(radar)\n",
    "                else:\n",
    "                    radar_out = torch.zeros(bs, 16).to(device)\n",
    "\n",
    "                if onoffdict['CAMERAS']:\n",
    "                    left_cam, center_cam, right_cam = process_imgs(gps, left_cam, center_cam, right_cam, crop_size = (150,150))\n",
    "                    lc_out = self.lc(left_cam)\n",
    "                    cc_out = self.cc(center_cam)\n",
    "                    rc_out = self.rc(right_cam)\n",
    "                else:\n",
    "                    lc_out = torch.zeros(bs, 16).to(device)\n",
    "                    cc_out = torch.zeros(bs, 16).to(device)\n",
    "                    rc_out = torch.zeros(bs, 16).to(device)\n",
    "\n",
    "                combined1 = torch.cat((gps_out, radar_out, lc_out, cc_out, rc_out), dim=1)\n",
    "\n",
    "                output = self.linear(combined1)\n",
    "                \n",
    "        \n",
    "        \n",
    "        \n",
    "        output = self.output_relu(output)\n",
    "        combined2 = torch.cat((output, Hdecoded.view(bs,-1)), dim=1)\n",
    "        output = self.output_fc1(combined2)\n",
    "        output = self.output_relu(output)\n",
    "        output = self.output_fc2(output)\n",
    "        output = output.view(bs, 1, num_H, num_H)\n",
    "\n",
    "        return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 76,
   "id": "5c1ed6df",
   "metadata": {},
   "outputs": [],
   "source": [
    "#complete task 2 model including encoder, decoder and channel\n",
    "class task2model(nn.Module):\n",
    "    def __init__(self, reduction=16):\n",
    "        super().__init__()\n",
    "        \n",
    "        self.en=torch.load(weight_pathb+\"task2Encoder.pth\")\n",
    "        self.allow_update = False  # Initially, do not allow weight updates\n",
    "        \n",
    "        self.de=Decoderwithmsi(reduction)\n",
    "        \n",
    "    \n",
    "   \n",
    "    def forward(self, Hin, gps, radar, left_cam, center_cam, right_cam, onoffdict): \n",
    "        \n",
    "        #Encoder\n",
    "        if self.allow_update:\n",
    "            Hencoded=self.en(Hin)\n",
    "        else:\n",
    "            with torch.no_grad():\n",
    "                Hencoded=self.en(Hin)\n",
    "        \n",
    "        \n",
    "        \n",
    "        \n",
    "        #Decoder\n",
    "        \n",
    "        Hdecoded=self.de(Hencoded,gps,radar,left_cam,center_cam,right_cam,onoffdict)\n",
    "        \n",
    "\n",
    "        return Hdecoded"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 77,
   "id": "eca2aed3",
   "metadata": {},
   "outputs": [],
   "source": [
    "model=task2model(reduction)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 78,
   "id": "7100468a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 79,
   "id": "812380a1",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Loss\n",
    "\n",
    "#criterion=nn.BCELoss()\n",
    "#criterion = nn.CrossEntropyLoss()\n",
    "criterion= nn.MSELoss().to(device)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "df8cbfba",
   "metadata": {},
   "source": [
    "# Inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 80,
   "id": "707e100e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Percentage Improvement Mean Achieved: 53.6652%\n",
      "Percentage Improvement Confidence Interval Achieved: 0.5160%\n",
      "Mean MSE: 0.5170\n",
      "95% Confidence Interval: (0.5112, 0.5228)\n",
      "Margin of Error: 0.0058\n"
     ]
    }
   ],
   "source": [
    "avg_mse_list = []\n",
    "improvement_list = []\n",
    "for _ in range(num_runs):\n",
    "    test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=True, num_workers=5)\n",
    "    model = torch.load(weight_path + \"task2.pth\").to(device)\n",
    "    avg_mse = run_test(model, test_loader, device, criterion)\n",
    "    avg_mse_list.append(avg_mse)\n",
    "    improvement = (target_loss.item() - avg_mse) / target_loss.item() * 100\n",
    "    improvement_list.append(improvement)\n",
    "mean_mse, margin_of_error = calculate_confidence_interval(avg_mse_list)\n",
    "improvement_mean, improvement_margin_of_error = calculate_confidence_interval(improvement_list)\n",
    "print(f'Percentage Improvement Mean Achieved: {improvement_mean:.4f}%')\n",
    "print(f'Percentage Improvement Confidence Interval Achieved: {improvement_margin_of_error:.4f}%')\n",
    "print(f\"Mean MSE: {mean_mse:.4f}\")\n",
    "print(f\"95% Confidence Interval: ({mean_mse - margin_of_error:.4f}, {mean_mse + margin_of_error:.4f})\")\n",
    "print(f\"Margin of Error: {margin_of_error:.4f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e1b94460",
   "metadata": {},
   "source": [
    "# Change Dictionary"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 81,
   "id": "50ee2c98",
   "metadata": {},
   "outputs": [],
   "source": [
    "onoffdict={'GPS': False, 'CAMERAS': False, 'RADAR': True}\n",
    "weight_path=f'models/TransNettask2/cr{reduction}/gps{onoffdict[\"GPS\"]}_cam{onoffdict[\"CAMERAS\"]}_rad{onoffdict[\"RADAR\"]}/'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 82,
   "id": "4fa7af9d",
   "metadata": {},
   "outputs": [],
   "source": [
    "if not os.path.exists(weight_path):\n",
    "    os.makedirs(weight_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 83,
   "id": "59e606cb",
   "metadata": {},
   "outputs": [],
   "source": [
    "task1_weight_path=f'models/task1/gps{onoffdict[\"GPS\"]}_cam{onoffdict[\"CAMERAS\"]}_rad{onoffdict[\"RADAR\"]}/'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 84,
   "id": "eddbcc05",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Decoderwithmsi(nn.Module):\n",
    "    def __init__(self, reduction):\n",
    "        super().__init__()\n",
    "        self.task1decoder = torch.load(task1_weight_path+\"task1Decoder.pth\")\n",
    "        self.gp = self.task1decoder.gp\n",
    "        self.rd = self.task1decoder.rd\n",
    "        self.lc = self.task1decoder.lc\n",
    "        self.cc = self.task1decoder.cc\n",
    "        self.rc = self.task1decoder.rc\n",
    "        self.bde = torch.load(weight_pathb+\"task2Decoder.pth\")\n",
    "        self.allow_update = False  # Initially, do not allow weight updates\n",
    "        #self.bde = baselinedecoder(reduction) \n",
    "\n",
    "        if int(num_H/2)*int(num_H/2) > 32:\n",
    "            self.linear = self.task1decoder.linear\n",
    "            self.output_fc1= nn.Linear(int(num_H/2)*int(num_H/2)+num_H*num_H, 2*num_H*num_H)\n",
    "            self.output_fc2 = nn.Linear(2*num_H*num_H, num_H*num_H)\n",
    "            self.output_relu = nn.ReLU()\n",
    "        else:\n",
    "            self.linear = self.task1decoder.linear\n",
    "            self.output_fc1 = nn.Linear(32+num_H*num_H, 2*num_H*num_H)\n",
    "            self.output_fc2 = nn.Linear(2*num_H*num_H, num_H*num_H)\n",
    "            self.output_relu = nn.ReLU()\n",
    "        \n",
    "        for m in self.modules():\n",
    "            if isinstance(m, (nn.Conv2d, nn.Linear)):\n",
    "                nn.init.xavier_uniform_(m.weight)\n",
    "            elif isinstance(m, nn.BatchNorm2d):\n",
    "                nn.init.constant_(m.weight, 1)\n",
    "                nn.init.constant_(m.bias, 0)\n",
    "    \n",
    "    def forward(self, Hencoded,gps,radar,left_cam,center_cam,right_cam,onoffdict):\n",
    "        bs = Hencoded.size(0)\n",
    "        if self.allow_update:\n",
    "            Hdecoded=self.bde(Hencoded)\n",
    "            if onoffdict['GPS']:\n",
    "                 gps_out = self.gp(gps)\n",
    "            else:\n",
    "                 gps_out = torch.zeros(bs, 16).to(device)\n",
    "\n",
    "            if onoffdict['RADAR']:\n",
    "                radar_out = self.rd(radar)\n",
    "            else:\n",
    "                radar_out = torch.zeros(bs, 16).to(device)\n",
    "\n",
    "            if onoffdict['CAMERAS']:\n",
    "                left_cam, center_cam, right_cam = process_imgs(gps, left_cam, center_cam, right_cam, crop_size = (150,150))\n",
    "                lc_out = self.lc(left_cam)\n",
    "                cc_out = self.cc(center_cam)\n",
    "                rc_out = self.rc(right_cam)\n",
    "            else:\n",
    "                lc_out = torch.zeros(bs, 16).to(device)\n",
    "                cc_out = torch.zeros(bs, 16).to(device)\n",
    "                rc_out = torch.zeros(bs, 16).to(device)\n",
    "\n",
    "            combined1 = torch.cat((gps_out, radar_out, lc_out, cc_out, rc_out), dim=1)\n",
    "\n",
    "            output = self.linear(combined1)\n",
    "        else:\n",
    "            with torch.no_grad():\n",
    "                Hdecoded=self.bde(Hencoded)\n",
    "                if onoffdict['GPS']:\n",
    "                     gps_out = self.gp(gps)\n",
    "                else:\n",
    "                     gps_out = torch.zeros(bs, 16).to(device)\n",
    "\n",
    "                if onoffdict['RADAR']:\n",
    "                    radar_out = self.rd(radar)\n",
    "                else:\n",
    "                    radar_out = torch.zeros(bs, 16).to(device)\n",
    "\n",
    "                if onoffdict['CAMERAS']:\n",
    "                    left_cam, center_cam, right_cam = process_imgs(gps, left_cam, center_cam, right_cam, crop_size = (150,150))\n",
    "                    lc_out = self.lc(left_cam)\n",
    "                    cc_out = self.cc(center_cam)\n",
    "                    rc_out = self.rc(right_cam)\n",
    "                else:\n",
    "                    lc_out = torch.zeros(bs, 16).to(device)\n",
    "                    cc_out = torch.zeros(bs, 16).to(device)\n",
    "                    rc_out = torch.zeros(bs, 16).to(device)\n",
    "\n",
    "                combined1 = torch.cat((gps_out, radar_out, lc_out, cc_out, rc_out), dim=1)\n",
    "\n",
    "                output = self.linear(combined1)\n",
    "                \n",
    "        \n",
    "        \n",
    "        \n",
    "        output = self.output_relu(output)\n",
    "        combined2 = torch.cat((output, Hdecoded.view(bs,-1)), dim=1)\n",
    "        output = self.output_fc1(combined2)\n",
    "        output = self.output_relu(output)\n",
    "        output = self.output_fc2(output)\n",
    "        output = output.view(bs, 1, num_H, num_H)\n",
    "\n",
    "        return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 85,
   "id": "13bad0f2",
   "metadata": {},
   "outputs": [],
   "source": [
    "#complete task 2 model including encoder, decoder and channel\n",
    "class task2model(nn.Module):\n",
    "    def __init__(self, reduction=16):\n",
    "        super().__init__()\n",
    "        \n",
    "        self.en=torch.load(weight_pathb+\"task2Encoder.pth\")\n",
    "        self.allow_update = False  # Initially, do not allow weight updates\n",
    "        \n",
    "        self.de=Decoderwithmsi(reduction)\n",
    "        \n",
    "    \n",
    "   \n",
    "    def forward(self, Hin, gps, radar, left_cam, center_cam, right_cam, onoffdict): \n",
    "        \n",
    "        #Encoder\n",
    "        if self.allow_update:\n",
    "            Hencoded=self.en(Hin)\n",
    "        else:\n",
    "            with torch.no_grad():\n",
    "                Hencoded=self.en(Hin)\n",
    "        \n",
    "        \n",
    "        \n",
    "        \n",
    "        #Decoder\n",
    "        \n",
    "        Hdecoded=self.de(Hencoded,gps,radar,left_cam,center_cam,right_cam,onoffdict)\n",
    "        \n",
    "\n",
    "        return Hdecoded"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 86,
   "id": "447150b0",
   "metadata": {},
   "outputs": [],
   "source": [
    "model=task2model(reduction)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 87,
   "id": "3727a922",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Training"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a9c4deac",
   "metadata": {},
   "source": [
    "# Inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 88,
   "id": "9049c0e3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Percentage Improvement Mean Achieved: 3.9752%\n",
      "Percentage Improvement Confidence Interval Achieved: 1.0468%\n",
      "Mean MSE: 1.0714\n",
      "95% Confidence Interval: (1.0597, 1.0831)\n",
      "Margin of Error: 0.0117\n"
     ]
    }
   ],
   "source": [
    "avg_mse_list = []\n",
    "improvement_list = []\n",
    "for _ in range(num_runs):\n",
    "    test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=True, num_workers=5)\n",
    "    model = torch.load(weight_path + \"task2.pth\").to(device)\n",
    "    avg_mse = run_test(model, test_loader, device, criterion)\n",
    "    avg_mse_list.append(avg_mse)\n",
    "    improvement = (target_loss.item() - avg_mse) / target_loss.item() * 100\n",
    "    improvement_list.append(improvement)\n",
    "mean_mse, margin_of_error = calculate_confidence_interval(avg_mse_list)\n",
    "improvement_mean, improvement_margin_of_error = calculate_confidence_interval(improvement_list)\n",
    "print(f'Percentage Improvement Mean Achieved: {improvement_mean:.4f}%')\n",
    "print(f'Percentage Improvement Confidence Interval Achieved: {improvement_margin_of_error:.4f}%')\n",
    "print(f\"Mean MSE: {mean_mse:.4f}\")\n",
    "print(f\"95% Confidence Interval: ({mean_mse - margin_of_error:.4f}, {mean_mse + margin_of_error:.4f})\")\n",
    "print(f\"Margin of Error: {margin_of_error:.4f}\")"
   ]
  }
 ],
 "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.10.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}