notebooks/simclr_tau_cifar10.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "3cbcc382",
   "metadata": {},
   "outputs": [],
   "source": [
    "data_root = '../datasets/cifar10/'\n",
    "ckp_root = '../checkpoints/cifar10/'\n",
    "code_root = '../'\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from matplotlib import colors\n",
    "from collections import Counter, defaultdict\n",
    "from numpy import linalg as LA\n",
    "from torchvision import transforms\n",
    "import torch\n",
    "from sklearn.svm import LinearSVC\n",
    "from torch.utils.data.sampler import SubsetRandomSampler\n",
    "import sys\n",
    "sys.path.append(code_root)\n",
    "from models.resnet import resnet18\n",
    "from utils import cvt_state_dict\n",
    "import torchvision\n",
    "from tqdm import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "7144b764-9731-4340-b6a4-56352dad0434",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Files already downloaded and verified\n"
     ]
    }
   ],
   "source": [
    "transform_train = transforms.Compose([\n",
    "    transforms.RandomCrop(32, padding=4),\n",
    "    transforms.RandomHorizontalFlip(),\n",
    "    transforms.ToTensor(),\n",
    "])\n",
    "transform_test = transforms.Compose([\n",
    "    transforms.ToTensor(),\n",
    "])\n",
    "\n",
    "\n",
    "train_datasets = torchvision.datasets.CIFAR10(root=data_root, \n",
    "                                              train=True, \n",
    "                                              download=True, \n",
    "                                              transform=transform_train)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "7de77ca7-b9a8-4e99-ad1b-f45303d0b09a",
   "metadata": {},
   "outputs": [],
   "source": [
    "val_splits = {\n",
    "1: 'cifar10_imbSub_with_subsets/split1_D_i.npy',\n",
    "2: 'cifar10_imbSub_with_subsets/split2_D_i.npy',\n",
    "3: 'cifar10_imbSub_with_subsets/split3_D_i.npy'\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "44fd750c-8fc1-46b8-9154-6532a386ea67",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Files already downloaded and verified\n"
     ]
    }
   ],
   "source": [
    "split = 1  # choose split to test [1,2,3]\n",
    "batch_size = 64\n",
    "train_idx = list(np.load(root + 'split/{}'.format(val_splits[split])))\n",
    "train_sampler = SubsetRandomSampler(train_idx)\n",
    "long_tail_loader = torch.utils.data.DataLoader(\n",
    "    train_datasets,\n",
    "    batch_size=batch_size, sampler=train_sampler, num_workers=0)\n",
    "    \n",
    "testset = torchvision.datasets.CIFAR10(root=data_root, train=False, download=True, transform=transform_test)\n",
    "test_loader = torch.utils.data.DataLoader(testset, batch_size=batch_size, shuffle=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "af58e921-0ebe-4328-a1da-605f5b01b9a3",
   "metadata": {},
   "outputs": [],
   "source": [
    "num_class = 10\n",
    "bnNameCnt = -1\n",
    "device = 'cuda'\n",
    "model = resnet18(pretrained=False, num_classes=num_class)\n",
    "model = model.to(device)\n",
    "\n",
    "path_checkpoint = root + f'models/simclr_TS_SP{split}.pt'\n",
    "checkpoint = torch.load(path_checkpoint, map_location=\"cpu\")\n",
    "if 'state_dict' in checkpoint:\n",
    "    state_dict = checkpoint['state_dict']\n",
    "elif 'P_state' in checkpoint:\n",
    "    state_dict = checkpoint['P_state']\n",
    "else:\n",
    "    state_dict = checkpoint\n",
    "\n",
    "try:\n",
    "    in_features = model.fc.in_features\n",
    "except AttributeError:\n",
    "    in_features = None\n",
    "state_dict = cvt_state_dict(state_dict, bnNameCnt, in_features, num_classes=num_class)\n",
    "\n",
    "model.load_state_dict(state_dict)\n",
    "# print('read checkpoint {}'.format(path_checkpoint))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "88d11d57-acc4-4c7a-a6f2-0b2405e72aef",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 175/175 [00:07<00:00, 22.03it/s]\n"
     ]
    }
   ],
   "source": [
    "# extract train long tail split features\n",
    "\n",
    "model.eval()\n",
    "train_LT_features = None\n",
    "train_LT_gt = []\n",
    "with torch.no_grad():\n",
    "    for data, target in tqdm(long_tail_loader):\n",
    "        data, target = data.to(device), target.to(device)\n",
    "        feats = model.eval()(data, features=True)\n",
    "        if train_LT_features is None:\n",
    "            train_LT_features = feats.detach().cpu()\n",
    "        else:\n",
    "            train_LT_features = torch.cat((train_LT_features, feats.detach().cpu()), dim=0)\n",
    "        train_LT_gt += target.cpu().numpy().tolist()\n",
    "\n",
    "data_normed = LA.norm(train_LT_features.numpy(), 2, axis=-1)\n",
    "train_LT_features = train_LT_features.numpy() / data_normed.reshape(-1, 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "10da7747-6bc5-4731-89bb-41ead91324cf",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 157/157 [00:05<00:00, 28.67it/s]\n"
     ]
    }
   ],
   "source": [
    "# extract test features\n",
    "\n",
    "model.eval()\n",
    "test_features = None\n",
    "test_gt = []\n",
    "with torch.no_grad():\n",
    "    for data, target in tqdm(test_loader):\n",
    "        data, target = data.to(device), target.to(device)\n",
    "        feats = model.eval()(data, features=True)\n",
    "        if test_features is None:\n",
    "            test_features = feats.detach().cpu()\n",
    "        else:\n",
    "            test_features = torch.cat((test_features, feats.detach().cpu()), dim=0)\n",
    "        test_gt += target.cpu().numpy().tolist()\n",
    "        \n",
    "data_normed = LA.norm(test_features.numpy(), 2, axis=-1)\n",
    "test_features = test_features.numpy() / data_normed.reshape(-1, 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "0612777e-6703-416b-8f0f-c7c6cdce577a",
   "metadata": {},
   "outputs": [],
   "source": [
    "dist_tmp = torch.cdist(torch.tensor(test_features), torch.tensor(train_LT_features))\n",
    "predicted = torch.argsort(dist_tmp, dim=1)\n",
    "sorted_dist_test_lt = predicted.numpy()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "f5e8e14b-56a7-4f1a-afec-3a3fc47af42b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SPLIT 1\n",
      "KNN@1  61.21\n",
      "KNN@10 61.20\n",
      "FS     71.49\n",
      "LT     62.92\n"
     ]
    }
   ],
   "source": [
    "print(f'SPLIT {split}')\n",
    "\n",
    "class_acc = []\n",
    "loc_gt = np.array(test_gt)\n",
    "loc_argsort = sorted_dist_test_lt\n",
    "train_loc_gt = np.array(train_LT_gt)\n",
    "for cl in range(num_class):\n",
    "    predictions = loc_argsort[loc_gt == cl, 0]\n",
    "    result = (train_loc_gt[predictions] == cl).sum() / (loc_gt == cl).sum()\n",
    "    class_acc.append(result)\n",
    "class_acc.append(np.mean(class_acc))\n",
    "print('KNN@1 ', f'{class_acc[-1]*100:.2f}')\n",
    "\n",
    "knn=10\n",
    "class_acc = []\n",
    "loc_gt = np.array(test_gt)\n",
    "loc_argsort = sorted_dist_test_lt\n",
    "train_loc_gt = np.array(train_LT_gt)\n",
    "for cl in range(num_class):\n",
    "    predict_mat = np.zeros(((loc_gt == cl).sum(), num_class))\n",
    "    predictions = loc_argsort[loc_gt == cl, :knn]\n",
    "    for cl2 in range(num_class):\n",
    "        result2 = (train_loc_gt[predictions] == cl2).sum(1)\n",
    "        predict_mat[:, cl2] = result2\n",
    "    result = (np.argmax(predict_mat, 1) == cl).sum() / (loc_gt == cl).sum()   \n",
    "    class_acc.append(result)\n",
    "class_acc.append(np.mean(class_acc))\n",
    "print('KNN@10', f'{class_acc[-1]*100:.2f}')\n",
    "\n",
    "local_data = train_LT_features\n",
    "local_gt = np.array(train_LT_gt)\n",
    "local_counter = Counter(local_gt)\n",
    "n_samples = local_counter.most_common()[-1][-1]\n",
    "np.random.seed(42)\n",
    "new_data = []\n",
    "new_gt = []\n",
    "for cl in range(num_class):\n",
    "    local_idxs = np.where(local_gt == cl)[0]\n",
    "    np.random.shuffle(local_idxs)\n",
    "    selected_idxs = local_idxs[:n_samples]\n",
    "    new_data.append(local_data[selected_idxs])\n",
    "    new_gt.extend([cl] * n_samples)\n",
    "\n",
    "new_data = np.concatenate(new_data, 0)\n",
    "fs_svm = LinearSVC(random_state=42)\n",
    "fs_svm.fit(new_data, new_gt)\n",
    "local_test = test_features\n",
    "local_gt_test = np.array(test_gt)\n",
    "predictions = fs_svm.predict(local_test)\n",
    "acc = (predictions == local_gt_test).sum() / len(local_gt_test)\n",
    "print('FS    ', f'{acc * 100:.2f}')\n",
    "\n",
    "\n",
    "lt_svm = LinearSVC(random_state=42)\n",
    "lt_svm.fit(local_data, local_gt)\n",
    "predictions = lt_svm.predict(local_test)\n",
    "acc = (predictions == local_gt_test).sum() / len(local_gt_test)\n",
    "print('LT    ', f'{acc * 100:.2f}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7e230d80-fbba-43a6-a79a-6f47e67c1ce6",
   "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.8.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}