UWB_Image_classification_v2.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "fFgWIgC39l0D",
    "tags": []
   },
   "outputs": [],
   "source": [
    "# !pip install gdown --upgrade\n",
    "# !pip install torchinfo\n",
    "# !pip install torchviz\n",
    "# !pip install wandb\n",
    "# !pip install -U imbalanced-learn\n",
    "# !pip install scikit-maad"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 35
    },
    "id": "-b5XdR_fQXUz",
    "outputId": "cfee7215-866a-4e74-a2fc-bda09f3870ff"
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "import pickle\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "import urllib\n",
    "import wandb\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torchvision.transforms as transforms\n",
    "\n",
    "from sklearn import preprocessing\n",
    "from torch.utils.data import Dataset\n",
    "from torch.utils.data import DataLoader\n",
    "from torchinfo import summary\n",
    "from tqdm.notebook import tqdm\n",
    "from scipy import signal\n",
    "\n",
    "torch.__version__ # 1.10.0+cu111"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "ZwD4PnvESYdK",
    "outputId": "ccb5b21a-2d46-45cd-953e-87d41f3412d7",
    "tags": []
   },
   "outputs": [],
   "source": [
    "!wandb login # 33d4ac1d40249bd82711d261c69583461dded7c7"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "wandb.login()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "P9T8UtrjN6WG"
   },
   "outputs": [],
   "source": [
    "# ! gdown --fuzzy https://drive.google.com/file/d/1iVah8T28-ib5YJUIyjI-sSbD7DiEfkuV/view?usp=share_link"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "b9pGUn-elJl_"
   },
   "source": [
    "# Preparing data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "FwKpcCNbRsGl",
    "tags": []
   },
   "outputs": [],
   "source": [
    "# !unzip uwb-pose-prediction.zip"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 424
    },
    "id": "ymveGtMwSpgi",
    "outputId": "81aa27f2-91e3-44ea-c033-d77e16b72d03",
    "tags": []
   },
   "outputs": [],
   "source": [
    "annot_df = pd.read_csv('annotations.csv')\n",
    "annot_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 269
    },
    "id": "iIDf0ultUTuc",
    "outputId": "27ec4fc6-d472-439f-ad80-1db8704f53fa",
    "tags": []
   },
   "outputs": [],
   "source": [
    "class_df = pd.read_csv('classes.csv')\n",
    "class_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "RbLIOQmpX-Za"
   },
   "outputs": [],
   "source": [
    "#Range time data\n",
    "def range_time(IQ_data, overlap = 1): \n",
    "    n_rd_history = 256\n",
    "    frame = []\n",
    "    frames = []\n",
    "    overlap_count = 0\n",
    "    check = 0\n",
    "    for iqini in IQ_data:\n",
    "        if len(frame)<n_rd_history:\n",
    "            frame.append(iqini)\n",
    "        else:  \n",
    "            if check == 0:\n",
    "                frame1 = np.array(frame)\n",
    "                frames.append(np.copy(frame1))\n",
    "                check = 1\n",
    "\n",
    "            if overlap_count < overlap:\n",
    "                frame.append(iqini)\n",
    "                overlap_count = overlap_count+1\n",
    "\n",
    "            if overlap_count == overlap:\n",
    "                frame = frame[overlap_count::]\n",
    "                frame1 = np.array(frame)\n",
    "                frames.append(np.copy(frame1))\n",
    "                overlap_count = 0\n",
    "    return np.stack(frames)\n",
    "\n",
    "import scipy.fftpack as fft\n",
    "#Range frequency data\n",
    "def range_frequency(datas, noise_threshold=None):\n",
    "    Range_frequency_frame = []\n",
    "    for data in datas:\n",
    "        jitter = 1e-10\n",
    "        # noise_threshold = -45\n",
    "        dB = True\n",
    "        #rd_history = np.hanning(n_rd_history)[:, None] * np.array(data)\n",
    "        # Range-Doppler\n",
    "        rd = fft.fft(data, axis=0)\n",
    "        rd = fft.fftshift(rd, axes=0)\n",
    "        rd = abs(rd)\n",
    "        if dB:\n",
    "            rd = 20 * np.log10(rd+jitter)\n",
    "            if noise_threshold is not None:\n",
    "                rd[rd < noise_threshold] = noise_threshold\n",
    "        Range_frequency_frame.append(rd)\n",
    "    return np.stack(Range_frequency_frame)\n",
    "\n",
    "def get_time_img(data, data_path = '/train/train/'):\n",
    "    # output = np.load(f'{data_path}{data_id}')\n",
    "    output = range_time(data, overlap=256)\n",
    "    return output.reshape(1,1,-1,56)\n",
    "    # return output.reshape(-1,56)\n",
    "\n",
    "def get_freq_img(data, noise_threshold=None):\n",
    "    # output = np.load(f'{data_path}{data_id}')\n",
    "    output = range_time(data, overlap=256)\n",
    "    # output = signal.detrend(output)\n",
    "    output = range_frequency(output, noise_threshold=noise_threshold)\n",
    "    return output.reshape(1,1,-1,56)\n",
    "    # return output.reshape(-1,56)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "qTKMTObXaVwa",
    "outputId": "3f6f1782-ba10-4f95-8d1b-4e76492a3124",
    "tags": []
   },
   "outputs": [],
   "source": [
    "# tmp = get_freq_img(np.load('train/train/001b0660-4c6e-4d07-8cd5-cd63578512f7.npy'), noise_threshold=-50)\n",
    "tmp = np.abs(get_time_img(np.load('train/train/001b0660-4c6e-4d07-8cd5-cd63578512f7.npy')))\n",
    "print(tmp.shape)\n",
    "tmp = torch.tensor(tmp.reshape(1,-1,56))\n",
    "resize = transforms.Resize((256,256))\n",
    "tmp = resize(tmp)\n",
    "plt.imshow(tmp.permute(1, 2, 0))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tmp = np.abs(get_time_img(np.load('train/train/001b0660-4c6e-4d07-8cd5-cd63578512f7.npy')))\n",
    "print(tmp.shape)\n",
    "tmp = torch.tensor(tmp.reshape(1,-1,56))\n",
    "resize = transforms.Resize((2560,2560))\n",
    "tmp = resize(tmp)\n",
    "resize = transforms.Resize((256,256))\n",
    "tmp = resize(tmp)\n",
    "plt.imshow(tmp.permute(1, 2, 0))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Tx18ZfToedJ7"
   },
   "outputs": [],
   "source": [
    "from pandas.core.common import random_state\n",
    "def train_test_split_df(df, test_size=0.1):\n",
    "    test_df = df.sample(frac = test_size, random_state = 42)\n",
    "    train_df = df.drop(test_df.index)\n",
    "    return train_df, test_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "tPkPxeHmgzUL",
    "outputId": "db1e4064-973d-4b83-8a9a-0ce808d4d656"
   },
   "outputs": [],
   "source": [
    "train_df, val_df = train_test_split_df(annot_df, test_size=0.2)\n",
    "print(train_df.shape, val_df.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 282
    },
    "id": "xtHTq66jg9nt",
    "outputId": "d877f9c4-947f-402b-a3ab-b4a2e94197a4"
   },
   "outputs": [],
   "source": [
    "plt.bar(train_df['class'].value_counts().index, height=train_df['class'].value_counts().values)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 282
    },
    "id": "ugF-wq7ThLP2",
    "outputId": "8e226f67-5d94-48bc-b472-63229d2d7bc3"
   },
   "outputs": [],
   "source": [
    "plt.bar(val_df['class'].value_counts().index, height=val_df['class'].value_counts().values)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def read_freq_data(df, main_path='train/train/', noise_threshold=None):\n",
    "    # features = np.array([], dtype=np.float32).reshape(1,2560,56)\n",
    "    # labels = np.array([], dtype=np.float32)\n",
    "    features = np.empty((0,1,2560,56))\n",
    "    labels = np.array([], dtype=np.float32)\n",
    "    for i in tqdm(range(len(df))):\n",
    "        feature_tmp = np.load(f\"{main_path}{df.iloc[i]['id']}.npy\")\n",
    "        if feature_tmp.shape != (2560,56):\n",
    "            continue\n",
    "        feature_tmp = get_freq_img(feature_tmp, noise_threshold=noise_threshold)\n",
    "        features = np.concatenate((features, feature_tmp), axis=0)\n",
    "        labels = np.concatenate((labels, df.iloc[i]['class'].reshape(1)), axis=0)\n",
    "\n",
    "    return features, labels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "N_Ulok3OUa7W"
   },
   "outputs": [],
   "source": [
    "def read_time_data(df, main_path='train/train/'):\n",
    "    # features = np.array([], dtype=np.float32).reshape(1,2560,56)\n",
    "    # labels = np.array([], dtype=np.float32)\n",
    "    features = np.empty((0,1,2560,56))\n",
    "    labels = np.array([], dtype=np.float32)\n",
    "    for i in tqdm(range(len(df))):\n",
    "        feature_tmp = np.load(f\"{main_path}{df.iloc[i]['id']}.npy\")\n",
    "        if feature_tmp.shape != (2560,56):\n",
    "            continue\n",
    "        feature_tmp = np.abs(get_time_img(feature_tmp))\n",
    "        features = np.concatenate((features, feature_tmp), axis=0)\n",
    "        labels = np.concatenate((labels, df.iloc[i]['class'].reshape(1)), axis=0)\n",
    "\n",
    "    return features, labels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "C7BZExcIFu0D"
   },
   "outputs": [],
   "source": [
    "class MyScaler:\n",
    "    def __init__(self, x_train, method='min-max'):\n",
    "        self.method = method\n",
    "        self.x_train = x_train\n",
    "        if method == 'min-max':\n",
    "            self.x_min = x_train.min()\n",
    "            self.x_max = x_train.max()\n",
    "        elif method == 'std':\n",
    "            self.x_mean = x_train.min()\n",
    "            self.x_std = x_train.std()\n",
    "        else:\n",
    "            raise NotImplementedError\n",
    "    \n",
    "    def transforms(self, x):\n",
    "        if self.method == 'min-max':\n",
    "            return (x - self.x_min)/(self.x_max-self.x_min)\n",
    "        elif self.method == 'std':\n",
    "            return (x - self.x_mean)/self.x_std"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "jux7vC5IbIZo",
    "outputId": "2fdee699-06c6-4fb1-95fd-c0523b4a1a23"
   },
   "outputs": [],
   "source": [
    "from imblearn.over_sampling import SMOTE\n",
    "\n",
    "noise_threshold = -50\n",
    "\n",
    "# x_train, y_train = read_freq_data(train_df, noise_threshold=noise_threshold)\n",
    "# x_val, y_val = read_freq_data(val_df, noise_threshold=noise_threshold)\n",
    "\n",
    "x_train, y_train = read_time_data(train_df)\n",
    "x_val, y_val = read_time_data(val_df)\n",
    "\n",
    "train_size = x_train.shape[0]\n",
    "sm = SMOTE(random_state=42)\n",
    "x_train, y_train = sm.fit_resample(x_train.reshape(train_size,-1),y_train)\n",
    "\n",
    "train_size = x_train.shape[0]\n",
    "x_train = x_train.reshape(train_size, 1, -1, 56)\n",
    "\n",
    "# x_scaler = MyScaler(x_train, method='std')\n",
    "# x_train = x_scaler.transforms(x_train)\n",
    "# x_val = x_scaler.transforms(x_val)\n",
    "\n",
    "print(x_train.shape, y_train.shape)\n",
    "print(x_val.shape, y_val.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "zB4mH7dXbVeJ"
   },
   "outputs": [],
   "source": [
    "class UWBDataset(Dataset):\n",
    "    def __init__(self, x, y):\n",
    "        self.x = x.astype(np.float32)\n",
    "        self.y = y.astype(np.float32)\n",
    "        print(self.x.shape)\n",
    "        print(self.y.shape)\n",
    "\n",
    "    def __getitem__(self, index):\n",
    "        x = self.x[index] # Retrieve feature data\n",
    "        resize = transforms.Resize((256,256))\n",
    "        x = resize(torch.tensor(x))\n",
    "        y = self.y[index] # Retrieve target\n",
    "        return x, y\n",
    "\n",
    "    def __len__(self):\n",
    "        return self.x.shape[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "vvqtHQPKiFAy",
    "outputId": "5c31dc90-9f33-4fd3-dcfb-3cabde9a63e8"
   },
   "outputs": [],
   "source": [
    "train_dataset = UWBDataset(x_train, y_train)\n",
    "val_dataset = UWBDataset(x_val, y_val)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 286
    },
    "id": "X1X0OKA5kOYo",
    "outputId": "938e009b-b921-4f2a-8e61-6468540925a7"
   },
   "outputs": [],
   "source": [
    "train_loader = DataLoader(train_dataset, batch_size=1, shuffle=True, pin_memory=True)\n",
    "for i in train_loader:\n",
    "    print(i[0].shape)\n",
    "    plt.imshow(i[0][0].permute(1, 2, 0))\n",
    "    break"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "wBk0qGHVlOB5"
   },
   "source": [
    "# Create model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "JihfCEL1oMQF",
    "tags": []
   },
   "outputs": [],
   "source": [
    "# !pip install timm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "kGdLWg-ToVOS"
   },
   "outputs": [],
   "source": [
    "# !pip install torchinfo"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "jpMfAzz7lOmc"
   },
   "outputs": [],
   "source": [
    "import timm\n",
    "import os \n",
    "from torchinfo import summary\n",
    "# os.mkdir('/kaggle/working/output')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "zyJVWWasn8mN"
   },
   "outputs": [],
   "source": [
    "def create_model(model_name, pretrained = True, num_classes=7):\n",
    "    model = timm.create_model(model_name, pretrained = pretrained, num_classes = num_classes)\n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "V8fze5kbn8pB"
   },
   "outputs": [],
   "source": [
    "def search_model(name):\n",
    "    for model_name in timm.list_models():\n",
    "        if name in model_name:\n",
    "            print(model_name)\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "PhhQIDrMn8rG",
    "tags": []
   },
   "outputs": [],
   "source": [
    "####################################################### List of model #######################################################\n",
    "name = 'efficientnetv2'\n",
    "search_model(name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "OhbKu6ohn8tq",
    "tags": []
   },
   "outputs": [],
   "source": [
    "model_name = 'efficientnetv2_s'\n",
    "\n",
    "model = create_model(model_name)\n",
    "model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "XzHgz8O8MV2z"
   },
   "outputs": [],
   "source": [
    "class MyModel(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(MyModel, self).__init__()\n",
    "        self.model = timm.create_model(model_name, pretrained=True, num_classes=7)\n",
    "        self.model.conv_stem = torch.nn.Conv2d(1, 24, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = self.model(x)\n",
    "        # x = nn.Softmax(dim=1)(x) # Softmax layer added to output\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "N3udVmetugzw",
    "outputId": "13a3d629-3da0-44d5-a57f-05baafff887c"
   },
   "outputs": [],
   "source": [
    "model = MyModel()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "GLofpqtln8wR",
    "outputId": "990d906b-b57e-4d01-ecf4-5f476b6ccacf"
   },
   "outputs": [],
   "source": [
    "summary(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 312
    },
    "id": "xoQa20sX0T-d",
    "outputId": "7ac071e0-298f-4aa2-dbe8-83a80f668700"
   },
   "outputs": [],
   "source": [
    "model = model.cuda()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "cPjnYFaHlhwi"
   },
   "source": [
    "# Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "c-QWne8-r3Q2",
    "outputId": "4868d588-a946-4fab-b891-6398b97651d2"
   },
   "outputs": [],
   "source": [
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "device"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "hqalHJTqxD8P"
   },
   "outputs": [],
   "source": [
    "run_name = f'{model_name}_lr0.01_range_time'\n",
    "result_path = f'results/{run_name}'\n",
    "if not os.path.exists(result_path):\n",
    "    os.mkdir(result_path)\n",
    "\n",
    "config = {\n",
    "    'architecture': model_name,\n",
    "    'lr': 0.01,\n",
    "    'scheduler_factor': 0.5,\n",
    "    'scheduler_patience': 3,\n",
    "    'early_stop_patience' : 10,\n",
    "    'scheduler_min_lr': 1e-4,\n",
    "    'epochs': 200,\n",
    "    'batch_size': 4,\n",
    "    'model_path': f'results/{run_name}/model.pth.tar',\n",
    "    'best_model_path': f'results/{run_name}/best_model.pth'\n",
    "}\n",
    "\n",
    "train_loader = DataLoader(train_dataset, batch_size=config['batch_size'], shuffle=True, pin_memory=True)\n",
    "val_loader = DataLoader(val_dataset, batch_size=config['batch_size'], shuffle=False, pin_memory=True)\n",
    "\n",
    "# Model\n",
    "model = model.to(device)\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=config['lr'])\n",
    "scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(\n",
    "    optimizer, \n",
    "    'min', \n",
    "    factor=config['scheduler_factor'], \n",
    "    patience=config['scheduler_patience'],\n",
    "    min_lr=config['scheduler_min_lr']\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 532,
     "referenced_widgets": [
      "4b8714156af34ba7afe13d6a8a2e6f21",
      "7b621352c66d44b1925c4ab0e64a7d6d",
      "8fba0349471f47e1821b991a3d85c9ae",
      "8765d9b4c7db4b6795f1668ab0fa31fc",
      "d6c1239d0a334e6f8bb7487b9b52d32f",
      "81515be516b4436f93ba38d1f38176fa",
      "4a2b4a7ea73e4e59bbd712e4220a3c33",
      "1962aaf2eda3460597e626d14342c6d2"
     ]
    },
    "id": "WIRrJIbpwfyw",
    "outputId": "5e32ac69-b71e-4559-b458-990b31ec83cd"
   },
   "outputs": [],
   "source": [
    "train_losses = []\n",
    "val_losses = []\n",
    "learning_rates = []\n",
    "loss_fn = nn.CrossEntropyLoss()\n",
    "early_stop_counter = 0\n",
    "best_val_loss = np.inf\n",
    "\n",
    "# Start wandb run\n",
    "wandb.init(\n",
    "    project=run_name,\n",
    "    config=config,\n",
    ")\n",
    "\n",
    "# Log parameters and gradients\n",
    "wandb.watch(model, log='all')\n",
    "\n",
    "for epoch in range(config['epochs']):  # loop over the dataset multiple times\n",
    "    \n",
    "    # Training\n",
    "    train_loss = []\n",
    "    current_lr = optimizer.param_groups[0]['lr']\n",
    "    learning_rates.append(current_lr)\n",
    "\n",
    "    # Flag model as training. Some layers behave differently in training and\n",
    "    # inference modes, such as dropout, BN, etc.\n",
    "    model.train()\n",
    "\n",
    "    print(f\"Training epoch {epoch+1}...\")\n",
    "    print(f\"Current LR: {current_lr}\")\n",
    "\n",
    "    for i, (inputs, y_true) in enumerate(tqdm(train_loader)):\n",
    "        # Transfer data from cpu to gpu\n",
    "        inputs = inputs.to(device)\n",
    "        y_true = y_true.to(device)\n",
    "        y_true = y_true.to(torch.int64)\n",
    "\n",
    "        # Reset the gradient\n",
    "        optimizer.zero_grad()\n",
    "\n",
    "        # Predict\n",
    "        y_pred = model(inputs)\n",
    "\n",
    "        # Calculate loss\n",
    "        loss = loss_fn(y_pred, y_true)\n",
    "\n",
    "        # Compute gradient\n",
    "        loss.backward()\n",
    "        \n",
    "        # Update parameters\n",
    "        optimizer.step()\n",
    "\n",
    "        # Log stuff\n",
    "        train_loss.append(loss)\n",
    "        \n",
    "    avg_train_loss = torch.stack(train_loss).mean().item()\n",
    "    train_losses.append(avg_train_loss)\n",
    "\n",
    "    print(f\"Epoch {epoch+1} train loss: {avg_train_loss:.4f}\")\n",
    "    \n",
    "    # Validation\n",
    "    model.eval()\n",
    "    with torch.no_grad(): # No gradient is required during validation\n",
    "        print(f\"Validating epoch {epoch+1}\")\n",
    "        val_loss = []\n",
    "        for i, (inputs, y_true) in enumerate(tqdm(val_loader)):\n",
    "            # Transfer data from cpu to gpu\n",
    "            inputs = inputs.to(device)\n",
    "            y_true = y_true.to(device)\n",
    "            y_true = y_true.to(torch.int64)\n",
    "            \n",
    "            # Predict\n",
    "            y_pred = model(inputs)\n",
    "\n",
    "            # Calculate loss\n",
    "            loss = loss_fn(y_pred, y_true)\n",
    "\n",
    "            # Log stuff\n",
    "            val_loss.append(loss)\n",
    "        \n",
    "        avg_val_loss = torch.stack(val_loss).mean().item()\n",
    "        val_losses.append(avg_val_loss)\n",
    "\n",
    "        print(f\"Epoch {epoch+1} val loss: {avg_val_loss:.4f}\")\n",
    "\n",
    "        # LR adjustment with scheduler\n",
    "        scheduler.step(avg_val_loss)\n",
    "\n",
    "        # Save checkpoint if val_loss is the best we got\n",
    "        best_val_loss = np.inf if epoch == 0 else min(val_losses[:-1])\n",
    "        if avg_val_loss < best_val_loss:\n",
    "            # Save whatever you want\n",
    "            state = {\n",
    "                'epoch': epoch,\n",
    "                'model': model.state_dict(),\n",
    "                'optimizer': optimizer.state_dict(),\n",
    "                'scheduler': scheduler.state_dict(),\n",
    "                'train_loss': avg_train_loss,\n",
    "                'val_loss': avg_val_loss,\n",
    "                'best_val_loss': best_val_loss,\n",
    "            }\n",
    "            \n",
    "            print(f\"Saving new best model..\")\n",
    "            torch.save(state, config['model_path'])\n",
    "\n",
    "    # Check for early stopping\n",
    "    if avg_val_loss < best_val_loss:\n",
    "        best_val_loss = avg_val_loss\n",
    "        early_stop_counter = 0\n",
    "        torch.save(model.state_dict(), config['best_model_path'])  # save the best model\n",
    "    else:\n",
    "        early_stop_counter += 1\n",
    "        if early_stop_counter == config['early_stop_patience']:\n",
    "            print(f'Early stopping at epoch {epoch}')\n",
    "            break\n",
    "    wandb.log({\n",
    "        'train_loss': avg_train_loss,\n",
    "        'val_loss': avg_val_loss,\n",
    "        'lr': current_lr,\n",
    "    })\n",
    "\n",
    "wandb.finish()\n",
    "print('Finished Training')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "IBYaf6Dyloax"
   },
   "source": [
    "# Evaluate model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 300
    },
    "id": "b1oTddic5Qei",
    "outputId": "fbe77302-eb8a-4c5c-86e0-cf6505c3db8b"
   },
   "outputs": [],
   "source": [
    "checkpoint = torch.load('')\n",
    "import copy\n",
    "best_model = copy.deepcopy(model)\n",
    "best_model.load_state_dict(checkpoint) # Load weights\n",
    "\n",
    "val_loader = DataLoader(val_dataset, batch_size=1, shuffle=False, pin_memory=True)\n",
    "pred_val_list = []\n",
    "for i in val_loader:\n",
    "    # print(i[0].shape)\n",
    "    # plt.imshow(i[0][0].permute(1, 2, 0))\n",
    "    best_model.eval()\n",
    "    with torch.no_grad():\n",
    "        pred = best_model(i[0].to(device))\n",
    "        pred = np.argmax(pred.to('cpu').numpy(), axis = 1).item()\n",
    "        pred_val_list.append(pred)\n",
    "\n",
    "from sklearn.metrics import classification_report, confusion_matrix, ConfusionMatrixDisplay\n",
    "\n",
    "print(classification_report(y_true=y_val, y_pred=pred_val_list))\n",
    "\n",
    "cm = confusion_matrix(y_true=y_val, y_pred=pred_val_list)\n",
    "print(cm)\n",
    "\n",
    "disp = ConfusionMatrixDisplay(confusion_matrix=cm)\n",
    "disp.plot()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "icBvVuOVls-F"
   },
   "source": [
    "# Run prediction for submission"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "pcz4kteNFBsN"
   },
   "outputs": [],
   "source": [
    "sample_sub = pd.read_csv('sample_submission.csv')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 206
    },
    "id": "tIZK4yL9EX_i",
    "outputId": "777a807c-fa50-4732-b949-04b7682b64a9"
   },
   "outputs": [],
   "source": [
    "sample_sub.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "kOHx5Rg7Ewbk"
   },
   "outputs": [],
   "source": [
    "test_data, _ = read_time_data(sample_sub, main_path = 'test/test/')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "dPHxohaTFhk9",
    "outputId": "42230492-5abf-4c25-aefa-6d060f259e95"
   },
   "outputs": [],
   "source": [
    "test_data = UWBDataset(test_data, _)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "7jbHuWW3Fr6I"
   },
   "outputs": [],
   "source": [
    "test_loader = DataLoader(test_data, batch_size=1, shuffle=False, pin_memory=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "checkpoint = torch.load('results/efficientnetv2_s_lr0.01_range_time/best_model.pth')\n",
    "import copy\n",
    "test_model = copy.deepcopy(model)\n",
    "test_model.load_state_dict(checkpoint) # Load weights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "sUBVRfSGGRCj",
    "tags": []
   },
   "outputs": [],
   "source": [
    "pred_list = []\n",
    "for i in test_loader:\n",
    "    # print(i[0].shape)\n",
    "    # plt.imshow(i[0][0].permute(1, 2, 0))\n",
    "    test_model.eval()\n",
    "    with torch.no_grad():\n",
    "        pred = test_model(i[0].to(device))\n",
    "        # print(pred)\n",
    "        # break\n",
    "        pred = np.argmax(pred.to('cpu').numpy(), axis = 1).item()\n",
    "        pred_list.append(pred)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "qhYLxFtZGdaP"
   },
   "outputs": [],
   "source": [
    "sample_sub['class'] = pred_list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 206
    },
    "id": "dqwPR__oH0gW",
    "outputId": "439e84d8-4a3b-4052-dad9-72db49e06fe8"
   },
   "outputs": [],
   "source": [
    "sample_sub"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "wpAn7mYIKRC1"
   },
   "outputs": [],
   "source": [
    "sample_sub.to_csv('results/efficientnetv2_s_lr0.01_range_time/kiddee3-range-time-3.csv', index = False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "provenance": []
  },
  "gpuClass": "standard",
  "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.9.13"
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "1962aaf2eda3460597e626d14342c6d2": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "ProgressStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "ProgressStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "bar_color": null,
      "description_width": ""
     }
    },
    "4a2b4a7ea73e4e59bbd712e4220a3c33": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "4b8714156af34ba7afe13d6a8a2e6f21": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "VBoxModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "VBoxModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "VBoxView",
      "box_style": "",
      "children": [
       "IPY_MODEL_7b621352c66d44b1925c4ab0e64a7d6d",
       "IPY_MODEL_8fba0349471f47e1821b991a3d85c9ae"
      ],
      "layout": "IPY_MODEL_8765d9b4c7db4b6795f1668ab0fa31fc"
     }
    },
    "7b621352c66d44b1925c4ab0e64a7d6d": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "LabelModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "LabelModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "LabelView",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_d6c1239d0a334e6f8bb7487b9b52d32f",
      "placeholder": "​",
      "style": "IPY_MODEL_81515be516b4436f93ba38d1f38176fa",
      "value": "Waiting for wandb.init()...\r"
     }
    },
    "81515be516b4436f93ba38d1f38176fa": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "DescriptionStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    },
    "8765d9b4c7db4b6795f1668ab0fa31fc": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "8fba0349471f47e1821b991a3d85c9ae": {
     "model_module": "@jupyter-widgets/controls",
     "model_module_version": "1.5.0",
     "model_name": "FloatProgressModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "FloatProgressModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "ProgressView",
      "bar_style": "",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_4a2b4a7ea73e4e59bbd712e4220a3c33",
      "max": 1,
      "min": 0,
      "orientation": "horizontal",
      "style": "IPY_MODEL_1962aaf2eda3460597e626d14342c6d2",
      "value": 0.01666950085000091
     }
    },
    "d6c1239d0a334e6f8bb7487b9b52d32f": {
     "model_module": "@jupyter-widgets/base",
     "model_module_version": "1.2.0",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    }
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}