linear_classifier.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Linear Classifier \n",
    "to be compared to `cosine similarity classifier `"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from sklearn import preprocessing\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "\n",
    "from torch import relu\n",
    "\n",
    "from torch.utils.data import DataLoader, Dataset\n",
    "\n",
    "import torchvision.datasets as dsets\n",
    "import torchvision.transforms as transforms"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "device(type='cuda', index=0)"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# GPU settings\n",
    "device = torch.device(0)\n",
    "device"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Few-shot Learning Settings\n",
    "K-shot N-ways"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "K = 10\n",
    "N = 1\n",
    "left_class = 7"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Get a dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_transform = transforms.Compose([transforms.ToTensor()])\n",
    "\n",
    "# import the `MNIST datasets`\n",
    "mnist_train = dsets.MNIST(root='data',\n",
    "                          train=True,\n",
    "                          transform=data_transform,\n",
    "                          download=True)\n",
    "\n",
    "mnist_test = dsets.MNIST(root='data',\n",
    "                          train=False,\n",
    "                          transform=data_transform,\n",
    "                          download=True)\n",
    "\n",
    "# build the `DataLoader`\n",
    "train_data_loader = DataLoader(mnist_train, batch_size=2**10)\n",
    "test_data_loader = DataLoader(mnist_test, batch_size=mnist_test.data.shape[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\Users\\Daesoo\\AppData\\Local\\Programs\\Python\\Python37\\venvs\\pytorch\\lib\\site-packages\\sklearn\\utils\\validation.py:63: DataConversionWarning: A column-vector y was passed when a 1d array was expected. Please change the shape of y to (n_samples, ), for example using ravel().\n",
      "  return f(*args, **kwargs)\n"
     ]
    }
   ],
   "source": [
    "# Label Encoder\n",
    "label_encoder  = preprocessing.LabelEncoder()\n",
    "\n",
    "targets = list(range(0, 10, 1))\n",
    "targets.pop(left_class)\n",
    "targets = np.array(targets).reshape(-1, 1)\n",
    "\n",
    "label_encoder.fit(targets);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([[[[0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          ...,\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.]]],\n",
      "\n",
      "\n",
      "        [[[0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          ...,\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.]]],\n",
      "\n",
      "\n",
      "        [[[0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          ...,\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.]]],\n",
      "\n",
      "\n",
      "        ...,\n",
      "\n",
      "\n",
      "        [[[0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          ...,\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.]]],\n",
      "\n",
      "\n",
      "        [[[0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          ...,\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.]]],\n",
      "\n",
      "\n",
      "        [[[0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          ...,\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.],\n",
      "          [0., 0., 0.,  ..., 0., 0., 0.]]]])\n",
      "tensor([5, 0, 4,  ..., 3, 7, 7])\n",
      "tensor(1.)\n",
      "torch.Size([1024, 1, 28, 28])\n"
     ]
    }
   ],
   "source": [
    "for x, y in train_data_loader:\n",
    "    print(x); print(y); print(x.max()); print(x.shape)\n",
    "    break"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Build a model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Model(nn.Module):\n",
    "    def __init__(self, in_size=28, embedding_feature_size=2, n_classes=10):\n",
    "        super().__init__()\n",
    "        \n",
    "        # Data properties\n",
    "        in_channels = 1\n",
    "        \n",
    "        # Define layers\n",
    "        # 1\n",
    "        self.conv1 = nn.Conv2d(in_channels, 12, kernel_size=3, padding=1)\n",
    "        self.bn1 = nn.BatchNorm2d(12)\n",
    "        self.maxpool1 = nn.MaxPool2d(kernel_size=2, stride=2)\n",
    "\n",
    "        # 2\n",
    "        self.conv2 = nn.Conv2d(12, 24, kernel_size=3, padding=1)\n",
    "        self.bn2 = nn.BatchNorm2d(24)\n",
    "        self.maxpool2 = nn.MaxPool2d(kernel_size=2, stride=2)\n",
    "\n",
    "        # 3\n",
    "        self.conv3 = nn.Conv2d(24, 48, kernel_size=3, padding=1)\n",
    "        self.bn3 = nn.BatchNorm2d(48)\n",
    "        self.maxpool3 = nn.MaxPool2d(kernel_size=2, stride=2)\n",
    "        \n",
    "        # 4\n",
    "        self.last_embedding_layer = nn.Linear(48*3*3, embedding_feature_size)\n",
    "        \n",
    "        # FC\n",
    "        self.fc1 = nn.Linear(embedding_feature_size, n_classes)\n",
    "        \n",
    "    def forward(self, x):\n",
    "        #============================\n",
    "        # Feature Extractor\n",
    "        h1 = self.maxpool1(relu(self.bn1(self.conv1(x))))\n",
    "        h2 = self.maxpool2(relu(self.bn2(self.conv2(h1))))\n",
    "        h3 = self.maxpool3(relu(self.bn3(self.conv3(h2))))\n",
    "        \n",
    "        batch_size = x.shape[0]\n",
    "        h3 = h3.view(batch_size, -1)  # 1\n",
    "        \n",
    "        z = self.last_embedding_layer(h3)\n",
    "        \n",
    "        # linear classifier\n",
    "        out = self.fc1(z)\n",
    "        \n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "model = Model(embedding_feature_size=64, n_classes=9).to(device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Compile"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)\n",
    "criterion = nn.CrossEntropyLoss()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "# settings\n",
    "n_epochs = 6"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "def leave_out_a_class(x, y, left_class):\n",
    "    \"leave out some class for the few-shot learning\"\n",
    "    indices = (y != left_class)\n",
    "    return x[indices, :, :, :], y[indices]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 0, loss: 0.546, test_acc: 0.973\n",
      "epoch: 1, loss: 0.084, test_acc: 0.982\n",
      "epoch: 2, loss: 0.056, test_acc: 0.986\n",
      "epoch: 3, loss: 0.043, test_acc: 0.989\n",
      "epoch: 4, loss: 0.034, test_acc: 0.989\n",
      "epoch: 5, loss: 0.028, test_acc: 0.989\n"
     ]
    }
   ],
   "source": [
    "train_hist = {\"epochs\": [], \"loss_per_epoch\": [], \"loss\": [], \"test_acc\": []}\n",
    "\n",
    "for epoch in range(n_epochs):\n",
    "    model.train()\n",
    "\n",
    "    loss_per_epoch = 0\n",
    "    iters = 0\n",
    "    for x, y in train_data_loader:\n",
    "        x, y = leave_out_a_class(x, y, left_class)\n",
    "        y = torch.tensor(label_encoder.transform(y))\n",
    "        x, y = x.to(device), y.to(device)\n",
    "        \n",
    "        optimizer.zero_grad()\n",
    "        \n",
    "        yhat = model(x)\n",
    "        \n",
    "        loss = criterion(yhat, y)\n",
    "        loss.backward()\n",
    "        \n",
    "        optimizer.step()\n",
    "        \n",
    "        # data storage\n",
    "        loss_per_epoch += loss\n",
    "        iters += 1\n",
    "        train_hist[\"loss\"].append(loss)\n",
    "        #print(round(loss.item(), 2), end=\" \")\n",
    "            \n",
    "    train_hist[\"epochs\"].append(epoch)\n",
    "    train_hist[\"loss_per_epoch\"].append(loss_per_epoch)\n",
    "    \n",
    "    # validation\n",
    "    with torch.no_grad():\n",
    "        model.eval()\n",
    "        \n",
    "        test_acc = 0\n",
    "        for x_test, y_test in test_data_loader:\n",
    "            x_test, y_test = leave_out_a_class(x_test, y_test, left_class)\n",
    "            y_test = torch.tensor(label_encoder.transform(y_test))\n",
    "            x_test, y_test = x_test.to(device), y_test.to(device)\n",
    "            yhat = torch.argmax(model(x_test.to(device)), axis=1)\n",
    "            test_acc += np.mean((yhat.to(\"cpu\") == y_test.to(\"cpu\")).numpy())\n",
    "        train_hist[\"test_acc\"].append(test_acc)\n",
    "    \n",
    "    print(\"epoch: {}, loss: {:0.3f}, test_acc: {:0.3f}\".format(epoch, loss_per_epoch/iters, test_acc))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Make `support set` and `query set`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "indices = (mnist_test.targets == left_class)\n",
    "\n",
    "class SupportSet(Dataset):\n",
    "    def __init__(self,):\n",
    "        super().__init__()\n",
    "        \n",
    "        support_set_X = []\n",
    "        support_set_Y = []\n",
    "        for c in range(0, 10):\n",
    "            indices = (mnist_test.targets == c)\n",
    "            support_set_X.append( mnist_test.data[indices, :, :][:N, :, :].numpy() / 255. )\n",
    "            support_set_Y.append( mnist_test.targets[indices][:N].numpy() )\n",
    "\n",
    "        self.support_set_X = torch.from_numpy(np.array(support_set_X)).view(N*10, 1, 28, 28)\n",
    "        self.support_set_Y = torch.from_numpy(np.array(support_set_Y).flatten())\n",
    "\n",
    "        self.len = self.support_set_X.shape[0]\n",
    "        \n",
    "    def __getitem__(self, idx):\n",
    "        return self.support_set_X[idx], self.support_set_Y[idx]\n",
    "    \n",
    "    def __len__(self):\n",
    "        return self.len\n",
    "\n",
    "support_set = SupportSet()\n",
    "support_set_data_loader = DataLoader(support_set, batch_size=2**10)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Freeze `Feature Extractor` and replace the last layer\n",
    "fine-tuning"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0 | ('conv1', Conv2d(1, 12, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)))\n",
      "1 | ('bn1', BatchNorm2d(12, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True))\n",
      "2 | ('maxpool1', MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False))\n",
      "3 | ('conv2', Conv2d(12, 24, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)))\n",
      "4 | ('bn2', BatchNorm2d(24, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True))\n",
      "5 | ('maxpool2', MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False))\n",
      "6 | ('conv3', Conv2d(24, 48, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)))\n",
      "7 | ('bn3', BatchNorm2d(48, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True))\n",
      "8 | ('maxpool3', MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False))\n",
      "9 | ('last_embedding_layer', Linear(in_features=432, out_features=64, bias=True))\n",
      "10 | ('fc1', Linear(in_features=64, out_features=9, bias=True))\n"
     ]
    }
   ],
   "source": [
    "# check out the layers with `.named_children()`\n",
    "for i, chd in enumerate(model.named_children()):\n",
    "    print(i, \"|\", chd)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "# freeze all the existing layers\n",
    "for param in model.parameters():\n",
    "    param.requires_grad = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "# replace the last layer\n",
    "model.fc1 = nn.Linear(64, 10).to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "conv1.weight | grad:False | device: cuda:0\n",
      "conv1.bias | grad:False | device: cuda:0\n",
      "bn1.weight | grad:False | device: cuda:0\n",
      "bn1.bias | grad:False | device: cuda:0\n",
      "conv2.weight | grad:False | device: cuda:0\n",
      "conv2.bias | grad:False | device: cuda:0\n",
      "bn2.weight | grad:False | device: cuda:0\n",
      "bn2.bias | grad:False | device: cuda:0\n",
      "conv3.weight | grad:False | device: cuda:0\n",
      "conv3.bias | grad:False | device: cuda:0\n",
      "bn3.weight | grad:False | device: cuda:0\n",
      "bn3.bias | grad:False | device: cuda:0\n",
      "last_embedding_layer.weight | grad:False | device: cuda:0\n",
      "last_embedding_layer.bias | grad:False | device: cuda:0\n",
      "fc1.weight | grad:True | device: cuda:0\n",
      "fc1.bias | grad:True | device: cuda:0\n"
     ]
    }
   ],
   "source": [
    "for name, param in model.named_parameters():\n",
    "    print(name, '| grad:{}'.format(param.requires_grad), '| device: {}'.format(param.device))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "optimizer = torch.optim.Adam([param for param in model.parameters() if param.requires_grad], \n",
    "                              lr=5e-4)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Fine-Tuning"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_hist = {\"epochs\": [], \"loss_per_epoch\": [], \"loss\": [], \"test_acc\": []}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_epochs = 200"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch: 0, loss: 3.406, test_acc: 0.060\n",
      "epoch: 10, loss: 2.789, test_acc: 0.088\n",
      "epoch: 20, loss: 2.244, test_acc: 0.159\n",
      "epoch: 30, loss: 1.783, test_acc: 0.357\n",
      "epoch: 40, loss: 1.406, test_acc: 0.554\n",
      "epoch: 50, loss: 1.107, test_acc: 0.685\n",
      "epoch: 60, loss: 0.874, test_acc: 0.747\n",
      "epoch: 70, loss: 0.696, test_acc: 0.795\n",
      "epoch: 80, loss: 0.560, test_acc: 0.844\n",
      "epoch: 90, loss: 0.457, test_acc: 0.875\n",
      "epoch: 100, loss: 0.379, test_acc: 0.894\n",
      "epoch: 110, loss: 0.319, test_acc: 0.902\n",
      "epoch: 120, loss: 0.272, test_acc: 0.908\n",
      "epoch: 130, loss: 0.236, test_acc: 0.911\n",
      "epoch: 140, loss: 0.207, test_acc: 0.913\n",
      "epoch: 150, loss: 0.183, test_acc: 0.916\n",
      "epoch: 160, loss: 0.164, test_acc: 0.917\n",
      "epoch: 170, loss: 0.148, test_acc: 0.919\n",
      "epoch: 180, loss: 0.134, test_acc: 0.920\n",
      "epoch: 190, loss: 0.122, test_acc: 0.920\n"
     ]
    }
   ],
   "source": [
    "for epoch in range(n_epochs):\n",
    "    model.train()\n",
    "\n",
    "    loss_per_epoch = 0\n",
    "    iters = 0\n",
    "    for x, y in support_set_data_loader:\n",
    "        x, y = x.to(device, dtype=torch.float), y.to(device)\n",
    "        \n",
    "        optimizer.zero_grad()\n",
    "        \n",
    "        yhat = model(x)\n",
    "        \n",
    "        loss = criterion(yhat, y)\n",
    "        loss.backward()\n",
    "        \n",
    "        optimizer.step()\n",
    "        \n",
    "        # data storage\n",
    "        loss_per_epoch += loss\n",
    "        iters += 1\n",
    "        train_hist[\"loss\"].append(loss)\n",
    "        #print(round(loss.item(), 2), end=\" \")\n",
    "            \n",
    "    train_hist[\"epochs\"].append(epoch)\n",
    "    train_hist[\"loss_per_epoch\"].append(loss_per_epoch)\n",
    "    \n",
    "    # validation\n",
    "    if (epoch % 10) == 0:\n",
    "        with torch.no_grad():\n",
    "            model.eval()\n",
    "            \n",
    "            test_acc = 0\n",
    "            for x_test, y_test in test_data_loader:\n",
    "                x_test, y_test = x_test.to(device, dtype=torch.float), y_test.to(device)\n",
    "                yhat = torch.argmax(model(x_test.to(device)), axis=1)\n",
    "                test_acc += np.mean((yhat.to(\"cpu\") == y_test.to(\"cpu\")).numpy())\n",
    "            train_hist[\"test_acc\"].append(test_acc)\n",
    "\n",
    "        print(\"epoch: {}, loss: {:0.3f}, test_acc: {:0.3f}\".format(epoch, loss_per_epoch/iters, test_acc))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Evaluate the learned representations"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`left_class` only"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "with torch.no_grad():\n",
    "    model.eval()\n",
    "    \n",
    "    indices = (mnist_test.targets == left_class)\n",
    "    \n",
    "    test_acc = 0\n",
    "    count = 0\n",
    "    for x_test, y_test in test_data_loader:\n",
    "        x_test, y_test = x_test.to(device), y_test.to(device)\n",
    "        x_test, y_test = x_test[indices,:,:], y_test[indices]\n",
    "        \n",
    "        yhat = torch.argmax(model(x_test.to(device)), axis=1)\n",
    "        \n",
    "        test_acc += np.mean((yhat.to(\"cpu\") == y_test.to(\"cpu\")).numpy())\n",
    "        count += 1\n",
    "        \n",
    "    test_acc /= count"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.6303501945525292"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "test_acc"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "over all classes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "class: 0 | acc: 0.978\n",
      "class: 1 | acc: 0.35\n",
      "class: 2 | acc: 0.905\n",
      "class: 3 | acc: 0.739\n",
      "class: 4 | acc: 0.829\n",
      "class: 5 | acc: 0.936\n",
      "class: 6 | acc: 0.817\n",
      "class: 7 | acc: 0.63\n",
      "class: 8 | acc: 0.896\n",
      "class: 9 | acc: 0.962\n",
      "\n",
      " overall acc: 0.804\n"
     ]
    }
   ],
   "source": [
    "with torch.no_grad():\n",
    "    model.eval()\n",
    "    \n",
    "    test_accs = []\n",
    "    for c in range(10):\n",
    "        indices = (mnist_test.targets == c)\n",
    "    \n",
    "        test_acc = 0\n",
    "        count = 0\n",
    "        for x_test, y_test in test_data_loader:\n",
    "            x_test, y_test = x_test.to(device), y_test.to(device)\n",
    "            x_test, y_test = x_test[indices,:,:], y_test[indices]\n",
    "\n",
    "            yhat = torch.argmax(model(x_test.to(device)), axis=1)\n",
    "\n",
    "            test_acc += np.mean((yhat.to(\"cpu\") == y_test.to(\"cpu\")).numpy())\n",
    "            count += 1\n",
    "\n",
    "        test_acc /= count\n",
    "        test_accs.append(test_acc)\n",
    "        print(f\"class: {c} | acc: {round(test_acc, 3)}\")\n",
    "\n",
    "print(\"\\n overall acc: {:0.3f}\".format(np.mean(test_accs)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pytorch",
   "language": "python",
   "name": "pytorch"
  },
  "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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}