v2.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "71c7a4d7",
   "metadata": {},
   "source": [
    "# Version 2: ConvNet + Cross Entropy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "4ebe396b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch, sys\n",
    "import argparse\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "import torch.utils.data.dataset\n",
    "from torch.utils.data import DataLoader\n",
    "from tqdm import tqdm\n",
    "\n",
    "import wandb\n",
    "from torch.utils.data import Dataset, DataLoader, Sampler\n",
    "from training_config import doodles, reals, doodle_size, real_size, NUM_CLASSES\n",
    "from utils import *  # bad practice, nvm\n",
    "from losses import compute_contrastive_loss_from_feats\n",
    "\n",
    "from torchvision import transforms\n",
    "import random\n",
    "import numpy as np\n",
    "import cv2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "da71673b",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "wandb.init(project=\"cs4243-project\", entity=\"rish-16\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "c12a9569",
   "metadata": {},
   "outputs": [],
   "source": [
    "def combined_dataset(datasets, size):\n",
    "    combined_dataset = {}\n",
    "    for name, dataset in datasets.items():\n",
    "        for class_name, class_data in dataset.items():\n",
    "            if class_name not in combined_dataset:\n",
    "                combined_dataset[class_name] = []\n",
    "            # resize data so they can be stacked\n",
    "            resized = []\n",
    "            for data in class_data:\n",
    "                resized.append(cv2.resize(data, (size, size), interpolation=cv2.INTER_AREA))\n",
    "            resized = np.stack(resized, axis=0)\n",
    "            combined_dataset[class_name].append(resized)\n",
    "    for class_name, lst_datasets in combined_dataset.items():\n",
    "        combined_dataset[class_name] = np.concatenate(lst_datasets, axis=0)\n",
    "    return combined_dataset\n",
    "\n",
    "\n",
    "class ImageDataset(Dataset):\n",
    "    DATASET_DIR = {True: 'dataset/dataset_train.npy', False: 'dataset/dataset_test.npy'}\n",
    "\n",
    "    def __init__(self, doodles_list, real_list, doodle_size, real_size, train: bool):\n",
    "        super(ImageDataset, self).__init__()\n",
    "\n",
    "        dataset = np.load(self.DATASET_DIR[train], allow_pickle=True)[()]\n",
    "\n",
    "        doodle_datasets = {name: data for name, data in dataset.items() if name in doodles_list}\n",
    "        real_datasets = {name: data for name, data in dataset.items() if name in real_list}\n",
    "        self.doodle_dict = combined_dataset(doodle_datasets, doodle_size)\n",
    "        self.real_dict = combined_dataset(real_datasets, real_size)\n",
    "\n",
    "        # sanity check\n",
    "        assert set(self.doodle_dict.keys()) == set(self.real_dict.keys()), \\\n",
    "            f'doodle and real images label classes do not match'\n",
    "\n",
    "        # process classes\n",
    "        label_idx = {}\n",
    "        for key in self.doodle_dict.keys():\n",
    "            if key not in label_idx:\n",
    "                label_idx[key] = len(label_idx)\n",
    "        self.label_idx = label_idx\n",
    "\n",
    "        # parse data and labels\n",
    "        self.doodle_data, self.doodle_label = self._return_x_y_pairs(self.doodle_dict, label_idx)\n",
    "        self.real_data, self.real_label = self._return_x_y_pairs(self.real_dict, label_idx)\n",
    "\n",
    "        # data preprocessing\n",
    "        self.doodle_preprocess = transforms.Compose([\n",
    "            transforms.ToPILImage(),\n",
    "            transforms.Resize(doodle_size),\n",
    "            transforms.ToTensor(),\n",
    "            transforms.Normalize((self.doodle_data/255).mean(), (self.doodle_data/255).std())   # IMPORTANT / 255\n",
    "        ])\n",
    "\n",
    "        self.real_preprocess = transforms.Compose([\n",
    "            transforms.ToPILImage(),\n",
    "            transforms.Resize(real_size),\n",
    "            transforms.ToTensor(),\n",
    "            transforms.Normalize((self.real_data/255).mean(axis=(0, 1, 2)), (self.real_data/255).std(axis=(0, 1, 2)))\n",
    "        ])\n",
    "\n",
    "        print(f'Train = {train}. Doodle list: {doodles_list}, \\n real list: {real_list}. \\n classes: {label_idx.keys()} \\n'\n",
    "              f'Doodle data size {len(self.doodle_data)}, real data size {len(self.real_data)}, '\n",
    "              f'ratio {len(self.doodle_data)/len(self.real_data)}')\n",
    "\n",
    "    def _return_x_y_pairs(self, data_dict, category_mapping):\n",
    "        xs, ys = [], []\n",
    "        for key in data_dict.keys():\n",
    "            data = data_dict[key]\n",
    "            labels = [category_mapping[key]] * len(data)\n",
    "            xs.append(data)\n",
    "            ys.extend(labels)\n",
    "        return np.concatenate(xs, axis=0), np.array(ys)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        # naive sampling scheme - sample with replacement\n",
    "        # sample label first so that doodle and real data belong to the same category\n",
    "        label = random.choice(list(self.label_idx.keys()))\n",
    "        doodle_data = self.doodle_preprocess(random.choice(self.doodle_dict[label]))\n",
    "        real_data = self.real_preprocess(random.choice(self.real_dict[label]))\n",
    "        numer_label = self.label_idx[label]\n",
    "        return doodle_data, numer_label, real_data, numer_label\n",
    "\n",
    "    def __len__(self):\n",
    "        return max(len(self.doodle_data), len(self.real_data)) # could be arbitrary number"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "909844bf",
   "metadata": {},
   "outputs": [],
   "source": [
    "class V2ConvNet(nn.Module):\n",
    "    def __init__(self, in_c, \n",
    "                 num_classes, \n",
    "                 channel_list=[64, 128, 192, 256, 512], \n",
    "                 pool_option=(1,1), \n",
    "                 hidden=256, \n",
    "                 dropout=0.2, \n",
    "                 add_layers=False):\n",
    "        super().__init__()\n",
    "        \n",
    "        layer1 = nn.Conv2d(in_c, channel_list[0], kernel_size=3)\n",
    "        layer2 = nn.Conv2d(channel_list[0], channel_list[0], kernel_size=3)\n",
    "        layers = [layer1, layer2]\n",
    "        \n",
    "        for i in range(1, len(channel_list)):\n",
    "            layers.append(\n",
    "                nn.Conv2d(channel_list[i-1], channel_list[i], kernel_size=3, stride=2, padding=1, bias=True)\n",
    "            )\n",
    "            layers.append(\n",
    "                nn.Conv2d(channel_list[i], channel_list[i], kernel_size=3, stride=2, padding=1, bias=True)\n",
    "            )\n",
    "            layers.append(\n",
    "                nn.BatchNorm2d(channel_list[i])\n",
    "            )\n",
    "            layers.append(\n",
    "                nn.Dropout(dropout)\n",
    "            )\n",
    "            layers.append(nn.ReLU())\n",
    "            \n",
    "        self.conv = nn.Sequential(*layers)\n",
    "        \n",
    "        self.flatten = nn.AdaptiveAvgPool2d(pool_option)\n",
    "            \n",
    "        self.fc = nn.Sequential(*[\n",
    "            nn.Linear(pool_option[0] * pool_option[1] * channel_list[-1], hidden),\n",
    "            nn.Linear(hidden, num_classes)\n",
    "        ])\n",
    "\n",
    "    def forward(self, x, return_feats=False):\n",
    "        feats = self.conv(x)\n",
    "        x = x.view(x.size(0), 512, -1).mean(2)\n",
    "        x = self.fc(x)\n",
    "\n",
    "        if return_feats:\n",
    "            return x, feats\n",
    "\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "b0c1e113",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "=================================================================\n",
      "Layer (type:depth-idx)                   Param #\n",
      "=================================================================\n",
      "V2ConvNet                                --\n",
      "├─Sequential: 1-1                        --\n",
      "│    └─Conv2d: 2-1                       896\n",
      "│    └─Conv2d: 2-2                       9,248\n",
      "│    └─Conv2d: 2-3                       36,992\n",
      "│    └─Conv2d: 2-4                       147,584\n",
      "│    └─BatchNorm2d: 2-5                  256\n",
      "│    └─Dropout: 2-6                      --\n",
      "│    └─ReLU: 2-7                         --\n",
      "│    └─Conv2d: 2-8                       590,336\n",
      "│    └─Conv2d: 2-9                       2,359,808\n",
      "│    └─BatchNorm2d: 2-10                 1,024\n",
      "│    └─Dropout: 2-11                     --\n",
      "│    └─ReLU: 2-12                        --\n",
      "├─AdaptiveAvgPool2d: 1-2                 --\n",
      "├─Sequential: 1-3                        --\n",
      "│    └─Linear: 2-13                      131,328\n",
      "│    └─Linear: 2-14                      2,313\n",
      "=================================================================\n",
      "Total params: 3,279,785\n",
      "Trainable params: 3,279,785\n",
      "Non-trainable params: 0\n",
      "=================================================================\n"
     ]
    }
   ],
   "source": [
    "from torchinfo import summary\n",
    "\n",
    "print (summary(V2ConvNet(3, 9, [32, 128, 512])))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "id": "9282f500",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([100, 9])\n"
     ]
    }
   ],
   "source": [
    "x = torch.rand(100, 3, 64, 64).cuda()\n",
    "net = V2ConvNet(3, 9, [64, 128, 192, 256, 512]).cuda()\n",
    "y = net(x)\n",
    "print (y.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 73,
   "id": "abbe7f41",
   "metadata": {},
   "outputs": [],
   "source": [
    "fix_seed(0)  # zero seed by default\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"1\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 74,
   "id": "2349123c",
   "metadata": {},
   "outputs": [],
   "source": [
    "sweep_config = {\n",
    "    'method': 'bayes'\n",
    "    }\n",
    "\n",
    "metric = {\n",
    "    'name': 'total_loss',\n",
    "    'goal': 'minimize'   \n",
    "    }\n",
    "\n",
    "sweep_config['metric'] = metric"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 75,
   "id": "3b79fe6b",
   "metadata": {},
   "outputs": [],
   "source": [
    "parameters_dict = {\n",
    "    'channel_combos': {\n",
    "        'values': [\n",
    "                [64, 128, 192, 256, 512],\n",
    "                [64, 128, 256, 512],\n",
    "                [32, 64, 64, 512],\n",
    "                [16, 32, 128, 512],\n",
    "                [32, 64, 192, 512],\n",
    "                [32, 128, 192, 512]\n",
    "            ]\n",
    "        },\n",
    "    'hidden_dim': {\n",
    "        \"values\": [64, 128, 256, 512]\n",
    "    },\n",
    "    'dropout': {\n",
    "          'values': [0.3, 0.4, 0.5]\n",
    "        }\n",
    "    }\n",
    "\n",
    "sweep_config['parameters'] = parameters_dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 76,
   "id": "5da83956",
   "metadata": {},
   "outputs": [],
   "source": [
    "parameters_dict.update({\n",
    "    'epochs': {\n",
    "        'value': 20\n",
    "    }\n",
    "})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 79,
   "id": "84f6e32b",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "parameters_dict.update({\n",
    "    'learning_rate': {\n",
    "        # a flat distribution between 0 and 0.1\n",
    "        'distribution': 'uniform',\n",
    "        'min': 0,\n",
    "        'max': 0.05\n",
    "    }\n",
    "})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 80,
   "id": "ec860240",
   "metadata": {},
   "outputs": [],
   "source": [
    "def build_model(doodle_channels, real_channels, hidden, dropout):\n",
    "    doodle_model = V2ConvNet(1, 9, channel_list=doodle_channels, hidden=hidden)\n",
    "    real_model = V2ConvNet(3, 9, channel_list=real_channels, hidden=hidden)\n",
    "    \n",
    "    return doodle_model, real_model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 81,
   "id": "f5454540",
   "metadata": {},
   "outputs": [],
   "source": [
    "def build_dataset(batch_size):\n",
    "    train_set = ImageDataset(doodles, reals, doodle_size, real_size, train=True)\n",
    "    val_set = ImageDataset(doodles, reals, doodle_size, real_size, train=False)\n",
    "\n",
    "    train_loader = DataLoader(train_set, batch_size=batch_size, shuffle=True, num_workers=16, pin_memory=True, drop_last=True)\n",
    "    val_loader = DataLoader(val_set, batch_size=batch_size, shuffle=False, num_workers=16, pin_memory=True, drop_last=True)\n",
    "\n",
    "    return train_loader, val_loader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 86,
   "id": "1932572a",
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_epoch(model1, model2, train_loader, val_loader, criterion, optimizer):\n",
    "    loss1_model1 = AverageMeter()\n",
    "    loss1_model2 = AverageMeter()\n",
    "    loss2_model1 = AverageMeter()\n",
    "    loss2_model2 = AverageMeter()\n",
    "    loss3_combined = AverageMeter()\n",
    "    acc_model1 = AverageMeter()\n",
    "    acc_model2 = AverageMeter()\n",
    "\n",
    "    model1.train()\n",
    "    model2.train()\n",
    "    total_loss = 0\n",
    "    \n",
    "    c1, c2, t = 0, 0, 0.1\n",
    "    \n",
    "    metadata = {}\n",
    "\n",
    "    for i, (x1, y1, x2, y2) in enumerate(train_loader):\n",
    "        # doodle, label, real, label\n",
    "        x1, y1, x2, y2 = x1.cuda(), y1.cuda(), x2.cuda(), y2.cuda()\n",
    "\n",
    "        # train model1 (doodle)\n",
    "        pred1, feats1 = model1(x1, return_feats=True)\n",
    "        loss_1 = criterion(pred1, y1)  # classification loss\n",
    "        loss1_model1.update(loss_1.item())\n",
    "        loss_model1 = loss_1\n",
    "\n",
    "        # train model2 (real)\n",
    "        pred2, feats2 = model2(x2, return_feats=True)\n",
    "        loss_1 = criterion(pred2, y2)  # classification loss\n",
    "        loss1_model2.update(loss_1.item())\n",
    "        loss_model2 = loss_1\n",
    "\n",
    "        loss = loss_model1 + loss_model2\n",
    "        total_loss += loss.item()\n",
    "\n",
    "        # statistics\n",
    "        acc_model1.update(compute_accuracy(pred1, y1))\n",
    "        acc_model2.update(compute_accuracy(pred2, y2))\n",
    "\n",
    "        # optimization\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        optimizer.zero_grad()\n",
    "\n",
    "    metadata.update({\n",
    "        'train acc 1': float(acc_model1.avg),\n",
    "        'train acc 2': float(acc_model2.avg),\n",
    "        'train epoch': epoch,\n",
    "        'l1m1': float(loss1_model1.avg),\n",
    "        'l1m2': float(loss1_model2.avg),\n",
    "        'total_loss': float(total_loss / base_bs)\n",
    "    })\n",
    "\n",
    "    # validation\n",
    "    model1.eval(), model1.eval()\n",
    "    acc_model1.reset(), acc_model2.reset()\n",
    "    with torch.no_grad():\n",
    "        for i, (x1, y1, x2, y2) in enumerate(val_loader):\n",
    "            pred1, feats1 = model1(x1, return_feats=True)\n",
    "            pred2, feats2 = model2(x2, return_feats=True)\n",
    "            acc_model1.update(compute_accuracy(pred1, y1))\n",
    "            acc_model2.update(compute_accuracy(pred2, y2))\n",
    "\n",
    "    metadata.update({\n",
    "        'val epoch': epoch,\n",
    "        'val acc 1': float(acc_model1.avg),\n",
    "        'val acc 2': float(acc_model2.avg),\n",
    "    })\n",
    "\n",
    "    scheduler.step()\n",
    "  \n",
    "    return metadata"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 87,
   "id": "db042349",
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_model(config=None):\n",
    "    num_epochs, base_bs, base_lr = 20, 512, 2e-2\n",
    "    \n",
    "    with wandb.init(config=config):\n",
    "        config = wandb.config\n",
    "        \n",
    "        model1, model2 = build_model(\n",
    "                            config.channel_combos,\n",
    "                            config.channel_combos,\n",
    "                            config.hidden_dim,\n",
    "                            config.dropout\n",
    "                        )\n",
    "        \n",
    "        model1 = model1.cuda()\n",
    "        model2 = model2.cuda()\n",
    "\n",
    "        train_loader, val_loader = build_dataset(base_bs)\n",
    "\n",
    "        optimizer = torch.optim.AdamW(params=list(model1.parameters()) + list(model2.parameters()), lr=config.learning_rate, weight_decay=3e-4)\n",
    "        criterion = nn.CrossEntropyLoss()\n",
    "        scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=num_epochs)\n",
    "\n",
    "        epoch_data = train_epoch(\n",
    "            model1, model2, \n",
    "            train_loader, \n",
    "            val_loader,\n",
    "            criterion,\n",
    "            optimizer\n",
    "        )\n",
    "        \n",
    "        wandb.log(epoch_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 88,
   "id": "9acf8aa5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Create sweep with ID: il4k0qg7\n",
      "Sweep URL: https://wandb.ai/rish-16/cs4243-project/sweeps/il4k0qg7\n"
     ]
    }
   ],
   "source": [
    "sweep_id = wandb.sweep(sweep_config, project=\"cs4243-project\", entity=\"rish-16\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2185a835",
   "metadata": {},
   "outputs": [],
   "source": [
    "wandb.agent(sweep_id, train_model, count=20)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "468eeb9d",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.7.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}