Fixmatch_Multiclass_Implementation.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Fixmatch_Multiclass_Implementation.ipynb",
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "code",
      "source": [
        "#@title\n",
        "!nvidia-smi -L"
      ],
      "metadata": {
        "id": "dwbHHlYUDeMG",
        "cellView": "form"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### configs"
      ],
      "metadata": {
        "id": "emY-znXfxHVk"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "NAME = \"Rifat\" #your name here\n",
        "PROJECT_NAME = \"Fixmatch_Multiclass_Implementation\"\n",
        "MODEL_TYPE = \"self_supervised\"\n",
        "ARCHITECTURE_NAME = \"tf_efficientnetv2_b0\"\n"
      ],
      "metadata": {
        "id": "LMeRqUabv7Yx"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Install dependency"
      ],
      "metadata": {
        "id": "y4S0T1lFxEt2"
      }
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Uq6OC6ZmADgB"
      },
      "outputs": [],
      "source": [
        "!pip install -q timm \n",
        "!pip install -q --upgrade --force-reinstall --no-deps kaggle\n",
        "!pip install -q opencv-python-headless==4.1.2.30  \n",
        "!pip install -q  --upgrade wandb "
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## **Download Kaggle Dataset and Dataframe create**"
      ],
      "metadata": {
        "id": "twYcPNxzxKhN"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import os\n",
        "import pandas as pd\n",
        "import numpy as np\n",
        "from imutils import paths\n",
        "from google.colab import files\n",
        "from sklearn.utils import shuffle\n",
        "from sklearn.model_selection import train_test_split"
      ],
      "metadata": {
        "id": "spkVEHm1mDfc"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "uploaded = files.upload()\n",
        "\n",
        "for fn in uploaded.keys():\n",
        "  print('User uploaded file \"{name}\" with length {length} bytes'.format(\n",
        "      name=fn, length=len(uploaded[fn])))\n",
        "  \n",
        "# Then move kaggle.json into the folder where the API expects to find it.\n",
        "!mkdir -p ~/.kaggle/ && mv kaggle.json ~/.kaggle/ && chmod 600 ~/.kaggle/kaggle.json"
      ],
      "metadata": {
        "id": "p9mnoKdZl330"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "**Dataset Download and unzip**"
      ],
      "metadata": {
        "id": "ggpxwBB6xNq7"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!kaggle datasets download -d trolukovich/apparel-images-dataset\n",
        "!mkdir dataset\n",
        "!unzip -q apparel-images-dataset.zip -d ./dataset"
      ],
      "metadata": {
        "id": "XjZ-2UFdl8f0"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-c1Mg6p7Ds91"
      },
      "source": [
        "**Dataframe Shuffle and Split**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "D-nArmYuDs92"
      },
      "outputs": [],
      "source": [
        "df_apparel = pd.read_csv(\"/content/df_apparel_multiclass.csv\")\n",
        "df = shuffle(df_apparel)\n",
        "ulb_dataframe , lb_dataframe = train_test_split(df,test_size = 0.2)\n",
        "lb_dataframe = shuffle(lb_dataframe)\n",
        "lb_dataframe, valid_dataframe = train_test_split(lb_dataframe,test_size = 0.2)\n",
        "print(len(ulb_dataframe))\n",
        "print(len(lb_dataframe))\n",
        "print(len(valid_dataframe))"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## **Useful imports**"
      ],
      "metadata": {
        "id": "rfYfBa6-a2Z_"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import pandas as pd\n",
        "import numpy as np\n",
        "import cv2\n",
        "\n",
        "import torch\n",
        "from torch import nn\n",
        "import torch.nn.functional as F\n",
        "from torch.utils.data import DataLoader, Dataset\n",
        "\n",
        "import torchvision\n",
        "from torchvision import transforms\n",
        "\n",
        "from tqdm.notebook import tqdm\n",
        "\n",
        "from sklearn.metrics import accuracy_score\n",
        "from timm.data.auto_augment import rand_augment_transform\n",
        "import PIL\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "import timm\n",
        "import time \n",
        "from collections import OrderedDict\n",
        "\n",
        "import wandb\n"
      ],
      "metadata": {
        "id": "hj0YwBeqB3on"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "**Variables**"
      ],
      "metadata": {
        "id": "oxbdHckhycGb"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# device is set to cuda if cuda is available\n",
        "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
        "mean = [0.485, 0.456, 0.406]\n",
        "std = [0.229, 0.224, 0.225]\n",
        "num_epochs = 10\n",
        "total_class = df_apparel[\"label\"].nunique()\n",
        "threshold = 0.90\n",
        "learning_rate = 0.001\n",
        "lb_to_ulb_ratio = 4\n",
        "batch_size_lb = 6\n",
        "save_path_checkpoints= \"/content/model/ckpts\"\n",
        "os.makedirs(save_path_checkpoints, exist_ok=True)"
      ],
      "metadata": {
        "id": "Sl_QPuYTmVy7"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "**Dataset Class**"
      ],
      "metadata": {
        "id": "_ud6TU5xyd5i"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "class ImageDataset(Dataset):\n",
        "    def __init__(self, dataframe,weak_transform,strong_transform,normalize,is_lb=False):\n",
        "        self.dataframe = dataframe\n",
        "        self.weak_transform = weak_transform\n",
        "        self.strong_transform = strong_transform\n",
        "        self.normalize = normalize\n",
        "        self.is_lb = is_lb\n",
        "        self.all_image_names = self.dataframe[:]['ImagePath']\n",
        "        self.all_image_label = self.dataframe.drop(['ImagePath'], axis=1)\n",
        "        \n",
        "    def __len__(self):\n",
        "        return len(self.all_image_names)\n",
        "\n",
        "    def __getitem__(self, index):\n",
        "        img_path = os.path.join(self.all_image_names.iloc[index])\n",
        "        image = PIL.Image.open(img_path)\n",
        "        if self.is_lb == True:\n",
        "          return self.normalize(image),int(self.all_image_label.iloc[index])\n",
        "        else:  \n",
        "          weak_image = self.weak_transform(image)\n",
        "          strong_image = self.strong_transform(image)\n",
        "          return self.normalize(weak_image),self.normalize(strong_image)\n",
        "         "
      ],
      "metadata": {
        "id": "Ctbcq1q2BJFu"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "**Dataloader and Transfrom**"
      ],
      "metadata": {
        "id": "zBODZ7_6yf1a"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def fun_transfrom():\n",
        "  strong_transform =  rand_augment_transform(\n",
        "    config_str='rand-m9-mstd0.5', \n",
        "    hparams={}\n",
        "              )\n",
        "\n",
        "  weak_transform = transforms.Compose([transforms.RandomHorizontalFlip()])\n",
        "      \n",
        "      \n",
        "  normalize = transforms.Compose([\n",
        "            transforms.Resize((224,224)),                      \n",
        "            transforms.ToTensor(),\n",
        "            transforms.Normalize(mean=mean, std=std)])\n",
        "\n",
        "\n",
        "  ulb_dataset = ImageDataset(\n",
        "      ulb_dataframe,\n",
        "      weak_transform,\n",
        "      strong_transform,\n",
        "      normalize,\n",
        "      is_lb=False\n",
        "      \n",
        "  )\n",
        "  lb_dataset = ImageDataset(\n",
        "      lb_dataframe,\n",
        "      weak_transform,\n",
        "      strong_transform,\n",
        "      normalize,\n",
        "      is_lb=True\n",
        "      \n",
        "  )\n",
        "\n",
        "  valid_dataset = ImageDataset(\n",
        "      valid_dataframe,\n",
        "      weak_transform,\n",
        "      strong_transform,\n",
        "      normalize,\n",
        "      is_lb=True\n",
        "      \n",
        "  )\n",
        "  \n",
        "  dataloader_ulb_dataset = DataLoader(ulb_dataset, batch_size=batch_size_lb * lb_to_ulb_ratio, shuffle=True, num_workers=2)\n",
        "  dataloader_lb_dataset = DataLoader(lb_dataset, batch_size=batch_size_lb, shuffle=True, num_workers=2)\n",
        "  dataloader_valid_dataset = DataLoader(valid_dataset, batch_size=20, shuffle=False, num_workers=2)\n",
        "\n",
        "  return dataloader_ulb_dataset,dataloader_lb_dataset,dataloader_valid_dataset\n",
        "\n",
        "\n",
        "dataloader_ulb_dataset,dataloader_lb_dataset,dataloader_valid_dataset = fun_transfrom()"
      ],
      "metadata": {
        "id": "38Xf_3C9ojx1"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jq1OJj0Fb7uV"
      },
      "source": [
        "**Wandb**"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1M97VqEEMzf_"
      },
      "outputs": [],
      "source": [
        "!wandb login "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "L3YKEACgfYzp"
      },
      "outputs": [],
      "source": [
        "class WandbLogger():\n",
        "    \"\"\"\n",
        "    This custom callback is used for logging training metrics to wandb for monitoring.\n",
        "    \n",
        "    \"\"\"\n",
        "    def __init__(self,project,entity,name,id,config,resume = \"allow\",):\n",
        "      self.project = project\n",
        "      self.entity = entity\n",
        "      self.name = name\n",
        "      self.id = id\n",
        "      self.config = config\n",
        "      self.resume = resume\n",
        "      wandb.init(project = self.project,entity = self.entity,\n",
        "                 name = self.name,id = self.id, \n",
        "                 config = self.config,resume = self.resume)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "X6pk20PrcAPR"
      },
      "outputs": [],
      "source": [
        "project= PROJECT_NAME\n",
        "entity=\"rakib1521\"\n",
        "\n",
        "\n",
        "name = f\"{PROJECT_NAME}_{ARCHITECTURE_NAME}\" #same name for multiple run is allowed but same id is not allowed\n",
        "id = f\"{PROJECT_NAME}_{ARCHITECTURE_NAME}\"\n",
        "\n",
        "wandb_config = {\"network\":ARCHITECTURE_NAME,\n",
        "                \"epoch\":num_epochs,\n",
        "                \"batch_size_lb\": batch_size_lb,\n",
        "                \"lb_to_ulb_ratio\":lb_to_ulb_ratio,\n",
        "                \"learning_rate\": learning_rate,\n",
        "                \"probability_threshold\": threshold,\n",
        "                }\n",
        "wandb_logger = WandbLogger(project,entity,name,id,wandb_config)    "
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Model Define"
      ],
      "metadata": {
        "id": "mV6ttB1Lym7u"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "model = timm.create_model('tf_efficientnetv2_b0')\n",
        "\n",
        "classifier = torch.nn.Sequential(OrderedDict([\n",
        "    ('fc1', torch.nn.Linear(model.classifier.in_features,total_class))\n",
        "]))\n",
        "\n",
        "model.classifier = classifier\n",
        "\n",
        "model = model.to(device)\n",
        "\n",
        "optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)"
      ],
      "metadata": {
        "id": "eFS7148AOhR3"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def train_fixMatch(threshold):\n",
        "    \n",
        "    model.train()\n",
        "    \n",
        "    losses = []\n",
        "    \n",
        "    pbar = tqdm(dataloader_lb_dataset)\n",
        "    \n",
        "    ulb_iter = iter(dataloader_ulb_dataset)\n",
        "    \n",
        "    for batch_lb in pbar:\n",
        "        \n",
        "        model.zero_grad()\n",
        "        \n",
        "        # labelled data\n",
        "        x_lb, y = batch_lb\n",
        "        x_lb, y = x_lb.to(device), y.to(device)\n",
        "        \n",
        "        # unlabelled data\n",
        "        try:\n",
        "            x_weak, x_strong = next(ulb_iter)\n",
        "        except StopIteration:\n",
        "            ulb_iter = iter(dataloader_ulb_dataset)\n",
        "            x_weak, x_strong = next(ulb_iter)\n",
        "                \n",
        "        x_weak, x_strong = x_weak.to(device), x_strong.to(device)\n",
        "        \n",
        "        # concat all x\n",
        "        all_x = torch.cat([x_lb, x_weak, x_strong], dim=0)\n",
        "        \n",
        "        # compute logits\n",
        "        all_logits = model(all_x)\n",
        "        \n",
        "        \n",
        "        # logits and loss for labelled data\n",
        "        logits_lb = all_logits[:x_lb.size(0)]\n",
        "\n",
        "        loss_lb = F.cross_entropy(logits_lb, y)\n",
        "\n",
        "        \n",
        "        # logits for unlabelled data\n",
        "        logits_ulb = all_logits[x_lb.size(0):]\n",
        "\n",
        "        logits_weak, logits_strong = torch.chunk(logits_ulb, 2, dim=0)\n",
        "\n",
        "        # stop gradient for weak augmented\n",
        "        logits_weak = logits_weak.detach() \n",
        "        \n",
        "        # compute class probailities\n",
        "        probs_weak = F.softmax(logits_weak, dim=1)\n",
        "\n",
        "        \n",
        "        # compute pseudo labels (torch.max outputs the maximum values and the argmax)\n",
        "        max_prob, pseudo_label = torch.max(probs_weak, dim=1)\n",
        "\n",
        "        \n",
        "        # mask for\n",
        "        mask = (max_prob > threshold).float() # [1, 0] [batch_size,]\n",
        "\n",
        "\n",
        "                \n",
        "        # mask non-confident prediction\n",
        "        pseudo_label = pseudo_label.masked_fill(mask == 0, 0) # [3, -1, ...]\n",
        "\n",
        "\n",
        "        \n",
        "        # unsupervised loss by ignoring non-confident prediction\n",
        "        loss_ulb = F.cross_entropy(logits_strong, pseudo_label, ignore_index=-1)\n",
        "\n",
        "\n",
        "\n",
        "        # total loss\n",
        "        loss = loss_lb + loss_ulb\n",
        "\n",
        "\n",
        "        \n",
        "        loss.backward()\n",
        "\n",
        "        optimizer.step()\n",
        "        \n",
        "        loss_item = loss.item()\n",
        "        \n",
        "        losses.append(loss_item)\n",
        "        \n",
        "        pbar.set_description(f'train loss = {np.array(losses).mean(): .3f}')\n",
        "        \n",
        "    return np.array(losses).mean()\n",
        "@torch.no_grad()\n",
        "def validate():\n",
        "    \n",
        "    model.eval()\n",
        "\n",
        "    labels_all = []\n",
        "    logits_all = []\n",
        "    \n",
        "    \n",
        "    for x, y in dataloader_valid_dataset:\n",
        "\n",
        "        x, y = x.to(device), y.to(device)\n",
        "\n",
        "        logits = model(x)\n",
        "        \n",
        "        labels_all += y.cpu().numpy().tolist()\n",
        "        logits_all += logits.cpu().numpy().tolist()\n",
        "     \n",
        "    prediction = np.argmax(np.array(logits_all), axis=-1)\n",
        "    \n",
        "    acc = accuracy_score(labels_all, prediction)\n",
        "                    \n",
        "    return acc"
      ],
      "metadata": {
        "id": "fMHARNntIN0g"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "losses = []\n",
        "accuracies = []\n",
        "for epoch in range(num_epochs):\n",
        "    train_loss = train_fixMatch(threshold=0.9)\n",
        "    val_acc= validate()\n",
        "    losses.append(train_loss)\n",
        "    accuracies.append(val_acc)\n",
        "    print(\"validation accuracy-{}\".format(val_acc))\n",
        "    wandb.log({ \"train_loss\": train_loss,\n",
        "                 \"valid_acc\":val_acc})  \n",
        "    filepath=f\"{save_path_checkpoints}/{PROJECT_NAME}_{MODEL_TYPE}-{ARCHITECTURE_NAME}-{epoch+1}_loss-{train_loss}.pt\"    \n",
        "    checkpoint= {\n",
        "                    \"epoch\" : epoch+1 ,\n",
        "                    \"model_weight\" : model.state_dict(),\n",
        "                    \"optimizer_state\" : optimizer.state_dict()\n",
        "                }\n",
        "    torch.save(checkpoint,filepath)\n",
        "    print(\"{} saved\".format(filepath)) \n",
        " "
      ],
      "metadata": {
        "id": "oWyJohbrInek"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "plt.plot(np.arange(num_epochs), losses)\n",
        "plt.title('traning Loss Vs epoch')\n",
        "plt.xlabel('epoch')\n",
        "plt.ylabel('traning Loss')\n",
        "plt.show()"
      ],
      "metadata": {
        "id": "RT_e4kSIotT9"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "plt.plot(np.arange(num_epochs), accuracies)\n",
        "plt.title('Validation accuracy Vs epoch')\n",
        "plt.xlabel('epoch')\n",
        "plt.ylabel('validation accuracy')\n",
        "plt.show()"
      ],
      "metadata": {
        "id": "7vv4K-cez3lk"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "@torch.no_grad()\n",
        "def validate():\n",
        "    \n",
        "    model.eval()\n",
        "\n",
        "    labels_all = []\n",
        "    logits_all = []\n",
        "    \n",
        "    \n",
        "    for x, y in dataloader_valid_dataset:\n",
        "\n",
        "        x, y = x.to(device), y.to(device)\n",
        "\n",
        "        logits = model(x)\n",
        "        \n",
        "        labels_all += y.cpu().numpy().tolist()\n",
        "        logits_all += logits.cpu().numpy().tolist()\n",
        "     \n",
        "    prediction = np.argmax(np.array(logits_all), axis=-1)\n",
        "    \n",
        "    acc = accuracy_score(labels_all, prediction)\n",
        "                    \n",
        "    return acc"
      ],
      "metadata": {
        "id": "tpSgf9y_cYQM"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import glob\n",
        "acc = []\n",
        "for num,path in enumerate(sorted(glob.glob(\"/content/model/ckpts/*\"))):\n",
        "  checkpoint = torch.load(path)\n",
        "  model.load_state_dict(checkpoint['model_weight'])\n",
        "  acc_valid  = validate()\n",
        "  acc.append(acc_valid)\n",
        "  #print(path)\n",
        "\n",
        "sorted(glob.glob(\"/content/model/ckpts/*\"))[np.argmax(acc)]  \n",
        "\n",
        "print(\"###\")\n",
        "model.load_state_dict(torch.load(sorted(glob.glob(\"/content/model/ckpts/*\"))[np.argmax(acc)])[\"model_weight\"])\n",
        "acc_valid  = validate()\n",
        "\n",
        "wandb.log({ \"Test_acc\":acc_valid})  \n",
        "\n",
        "wandb.finish()"
      ],
      "metadata": {
        "id": "LaEzUbwAyA-I"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}