project.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set up the environment \n",
    "import torch\n",
    "\n",
    "device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')\n",
    "\n",
    "compress_out_dim = 256"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Clone the ResNet implementation we are going to use\n",
    "# state_dict_path = 'resnet_20_64_pretrain.pt'\n",
    "state_dict_path = 'resnet_20_{}_pretrain.pt'.format(compress_out_dim)\n",
    "import os\n",
    "import shutil\n",
    "\n",
    "if not os.path.exists(state_dict_path):\n",
    "    !git clone https://github.com/geos98/CS294-082.git\n",
    "    shutil.copy(os.path.join('CS294-082', state_dict_path), '.')\n",
    "# !git clone https://github.com/akamaster/pytorch_resnet_cifar10"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<All keys matched successfully>\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "ResNet(\n",
       "  (conv1): Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "  (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "  (layer1): Sequential(\n",
       "    (0): BasicBlock(\n",
       "      (conv1): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (shortcut): Sequential()\n",
       "    )\n",
       "    (1): BasicBlock(\n",
       "      (conv1): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (shortcut): Sequential()\n",
       "    )\n",
       "    (2): BasicBlock(\n",
       "      (conv1): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (shortcut): Sequential()\n",
       "    )\n",
       "  )\n",
       "  (layer2): Sequential(\n",
       "    (0): BasicBlock(\n",
       "      (conv1): Conv2d(64, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
       "      (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (shortcut): LambdaLayer()\n",
       "    )\n",
       "    (1): BasicBlock(\n",
       "      (conv1): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (shortcut): Sequential()\n",
       "    )\n",
       "    (2): BasicBlock(\n",
       "      (conv1): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv2): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (shortcut): Sequential()\n",
       "    )\n",
       "  )\n",
       "  (layer3): Sequential(\n",
       "    (0): BasicBlock(\n",
       "      (conv1): Conv2d(128, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
       "      (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (shortcut): LambdaLayer()\n",
       "    )\n",
       "    (1): BasicBlock(\n",
       "      (conv1): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (shortcut): Sequential()\n",
       "    )\n",
       "    (2): BasicBlock(\n",
       "      (conv1): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (conv2): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "      (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      (shortcut): Sequential()\n",
       "    )\n",
       "  )\n",
       "  (linear): Linear(in_features=256, out_features=10, bias=True)\n",
       ")"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "if compress_out_dim == 256:\n",
    "    from resnet_256_out import resnet20\n",
    "if compress_out_dim == 64:\n",
    "    from resnet_64_out import resnet20\n",
    "\n",
    "state_dict = torch.load(state_dict_path, map_location=device)\n",
    "\n",
    "model = resnet20()\n",
    "miss = model.load_state_dict(state_dict)\n",
    "\n",
    "print(miss)\n",
    "\n",
    "model.to(device)\n",
    "model.eval()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Files already downloaded and verified\n",
      "Files already downloaded and verified\n",
      "Files already downloaded and verified\n"
     ]
    }
   ],
   "source": [
    "# Create datasets\n",
    "import torchvision.transforms as transforms\n",
    "import torchvision\n",
    "\n",
    "from torch.utils.data import DataLoader\n",
    "\n",
    "normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],\n",
    "                                 std=[0.229, 0.224, 0.225])\n",
    "\n",
    "train_transform = transforms.Compose([\n",
    "    transforms.RandomHorizontalFlip(),\n",
    "    transforms.RandomCrop(32, 4),\n",
    "    transforms.ToTensor(),\n",
    "    normalize,\n",
    "])\n",
    "\n",
    "test_transform = transforms.Compose([\n",
    "    transforms.ToTensor(),\n",
    "    normalize,\n",
    "])\n",
    "\n",
    "trainset = torchvision.datasets.CIFAR10(root='./data', train=True,\n",
    "                                        download=True, transform=train_transform)\n",
    "\n",
    "# Train data with test set trainsform (i.e., no crop, no flip)\n",
    "traintestset = torchvision.datasets.CIFAR10(root='./data', train=True,\n",
    "                                        download=True, transform=test_transform)\n",
    "testset = torchvision.datasets.CIFAR10(root='./data', train=False,\n",
    "                                        download=True, transform=test_transform)\n",
    "\n",
    "train_loader = DataLoader(trainset, batch_size=128, shuffle=True, pin_memory=True)\n",
    "test_loader = DataLoader(testset, batch_size=128, shuffle=False, pin_memory=True)\n",
    "traintest_loader = DataLoader(traintestset, batch_size=128, shuffle=False, pin_memory=True)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 79/79 [01:11<00:00,  1.10it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Testing Set: Predicted 9410 instances out of 10000\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 391/391 [06:03<00:00,  1.07it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Trainning Set: Predicted 50000 instances out of 50000\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "# Evaluation on accuracy\n",
    "from tqdm import tqdm\n",
    "\n",
    "\n",
    "with torch.no_grad():\n",
    "    model.eval()\n",
    "\n",
    "    correct_test = 0\n",
    "    total_test = 0\n",
    "    # Accuracy on test set\n",
    "    for p, l in tqdm(test_loader):\n",
    "        p = p.to(device)\n",
    "        out = model(p)\n",
    "        _, predicted = out.max(1)\n",
    "\n",
    "        correct_test += predicted.eq(l).sum().item()\n",
    "        total_test += l.size(0)\n",
    "    print('Testing Set: Predicted {} instances out of {}'.format(correct_test, total_test))\n",
    "\n",
    "    correct_train = 0\n",
    "    total_train = 0\n",
    "    # Accuracy on test set\n",
    "    for p, l in tqdm(traintest_loader):\n",
    "        p = p.to(device)\n",
    "        out = model(p)\n",
    "        _, predicted = out.max(1)\n",
    "\n",
    "        correct_train += predicted.eq(l).sum().item()\n",
    "        total_train += l.size(0)\n",
    "    print('Trainning Set: Predicted {} instances out of {}'.format(correct_train, total_train))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Q1: What is the variable the machine learner is supposed to predict? How accurate is the labeling? What is the annotator agreement (measured)?\n",
    "The CIFAR-10 dataset consists of 60,000 images of 10 classes which are  `airplane`, `automobile`, `bird`, `cat`, `deer`, `dog`, `frog`, `horse`, `ship` and `truck`. The machine learner should take-in an image and predict what is the object in the image by outputing one of the 10 classes listed above.\n",
    "\n",
    "There are a total of 22 known and validated label error exists in this dataset[1]. The label error was found algorithmically and further verified by human. This is 99.96% label accuracy. For the purpose of this project, we will ignore this error and assume the labelling is correct. \n",
    "\n",
    "Researchers at the University of Toronto paid human labler to label the dataset[2]. Each labelling is done by one human labler and verified by the one of the researchers after initial labelling. The in-accuracy could be caused by human error due to the large workload needed to label this dataset."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Q2: What is the required accuracy metric for success? How much data do we have to train the prediction of the variable? Are the classes balanced? How many modalities could be exploited in the data? Is there temporal information? How much noise are we expecting? Do we expect bias?\n",
    "\n",
    "Since this dataset serves as a benchmark towards the machine learner's ability to classify subjects in the image, the required accuracy metric is higher the better. However, there were no real-world implication for the success at this time.\n",
    "\n",
    "We have 60,000 images, splited into a trainning set of 50,000 images and testing set of 10,000 images. \n",
    "\n",
    "The classes are perfectly balanced in both trainning set and testing set. Meaning that in total there are 6,000 images of each class.\n",
    "\n",
    "There is only one modality, the visual, that we can exploit. \n",
    "\n",
    "There are no temporal information. \n",
    "\n",
    "Since every sample is an image with a background, and we are only interested in the main subject, this means that all the background portion of the image is effectively noise for our purpose. Therefore, we need to apply strategies like convolution to extract important features.\n",
    "\n",
    "We expect the machine learner to focus on the center of the image since that is the place where most subject would exists in the image."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Q3: What is the MEC of the data? What is the expected MEC for a neural net?\n",
    "\n",
    "Since we only care about the MEC of the data after the \"compression\", we will create a CSV file of data after compression"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 60000/60000 [07:25<00:00, 134.58it/s]\n"
     ]
    }
   ],
   "source": [
    "# Generate .csv from pre-trained ResNet-18 (only using the conv layer to extract deepfeature)\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "from tqdm import tqdm\n",
    "from itertools import chain\n",
    "\n",
    "# Get pre-trained conv output\n",
    "def run_conv(x):\n",
    "    out = F.relu(model.bn1(model.conv1(x)))\n",
    "    out = model.layer1(out)\n",
    "    out = model.layer2(out)\n",
    "    out = model.layer3(out)\n",
    "    out = F.avg_pool2d(out, out.size()[3])\n",
    "    out = out.view(out.size(0), -1)\n",
    "\n",
    "    return out\n",
    "\n",
    "compressed_file = 'ompressed_{}.csv'.format(compress_out_dim)\n",
    "with open(compressed_file, 'w+') as fout:\n",
    "    with torch.no_grad():\n",
    "        model.eval()\n",
    "        total_len = len(traintestset) + len(testset)\n",
    "        for p, l in tqdm(chain(traintestset, testset), total = total_len):\n",
    "            p = torch.unsqueeze(p, 0).to(device)\n",
    "            conv_data = run_conv(p)\n",
    "            np_data = conv_data.flatten().cpu().numpy()\n",
    "            CSV_str = ','.join(['%.15f' % num for num in np_data])\n",
    "            CSV_str += ',{}'.format(l)\n",
    "            fout.write(CSV_str + '\\n')\n",
    "            "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can either run Brainome on the dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Install Brainome for analysis\n",
    "#%pip install brainome\n",
    "#!brainome login\n",
    "#!brainome compressed.csv -y -headerless -f NN -vvv"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Or we can run our own analysis"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "dim of data:  256\n",
      "MEC =  4096.156322646796  bits\n"
     ]
    }
   ],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "import math\n",
    "\n",
    "from IPython.display import display\n",
    "\n",
    "def memorize(df: pd.DataFrame, drop = 0):\n",
    "    dfs = pd.DataFrame({'sum': df.iloc[:, :-1].sum(axis = 1), 'label': df.iloc[:, 64]})\n",
    "\n",
    "    drop_count = int(drop * df.shape[0])\n",
    "    drops = np.random.choice(dfs.index, drop_count, replace = False)\n",
    "    dfs = dfs.drop(drops)\n",
    "\n",
    "    dfs.sort_values('sum')\n",
    "    thr = (dfs['label'].diff() != 0).sum()\n",
    "    mec = math.log2((thr + 1))\n",
    "\n",
    "    return mec\n",
    "\n",
    "# Implement Alogrithm 4 in the book (page 93)\n",
    "df = pd.read_csv(compressed_file, header=None)\n",
    "dim = df.shape[1] - 1\n",
    "\n",
    "print('dim of data: ', dim)\n",
    "\n",
    "mec_mem = memorize(df)\n",
    "mec = (mec_mem * (dim + 1)) + mec_mem + 1\n",
    "print('MEC = ', mec, ' bits')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As calculated above, we find that the MEC of the data is around 15.72 bits.\n",
    "\n",
    "\n",
    "The decision layer in ResNet-20 is a simple fully connected layer with 64 input and 10 output. That is, 10 neurons each with 64 input.Now, each individual neuron therefore have parameter count of 64 + 1 = 65 and the MEC of the network is  \n",
    "65 * 10 = 650 bits"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Q4: What is the expected generalization in bits/bit and as a consequence the average resilience in dB? Is the resilience enough for the task? How bad can adversarial examples be? Do we expect data drift?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 60000/60000 [00:00<00:00, 93908.46it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Correctly predicted instances:  59410\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "# Test the data on pre-trained decision layer\n",
    "from torch.utils.data.dataset import Dataset\n",
    "class CompressedCifar10(Dataset):\n",
    "    def __init__(self, type = 'all'):\n",
    "        df = pd.read_csv(compressed_file, header=None)\n",
    "\n",
    "        row_begin = 0\n",
    "        row_end = 60000\n",
    "\n",
    "        if type == 'train':\n",
    "            row_end = 50000\n",
    "        if type == 'test':\n",
    "            row_begin = 50000\n",
    "\n",
    "        col_max = df.shape[1]\n",
    "        \n",
    "        x = df.iloc[row_begin:row_end, 0:col_max -  1].values\n",
    "        y = df.iloc[row_begin:row_end, col_max - 1].values\n",
    "\n",
    "        self.x_data = torch.tensor(x, dtype=torch.float32, device=device)\n",
    "        self.y_data = torch.tensor(y, dtype=torch.long, device=device)\n",
    "    \n",
    "    def __len__(self):\n",
    "        return len(self.y_data)\n",
    "    \n",
    "    def __getitem__(self, index):\n",
    "        return self.x_data[index], self.y_data[index]\n",
    "\n",
    "total_correct = 0\n",
    "with torch.no_grad():\n",
    "    dataset = CompressedCifar10()\n",
    "\n",
    "    for x, y in tqdm(dataset):\n",
    "        preidct = model.linear(x)\n",
    "        predict = torch.argmax(preidct)\n",
    "\n",
    "        if predict == y:\n",
    "            total_correct += 1\n",
    "\n",
    "print('Correctly predicted instances: ', total_correct)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Therefore, combine MEC = 650 bit, we get the Generalization is\n",
    "\n",
    "$$G = \\frac{59038  \\log_2{10}}{MEC} = 188.81 \\frac{\\text{bits}}{\\text{bit}}$$\n",
    "\n",
    "Which would then traslate to resilience $R = 20 * \\log_{10}G = 20 * \\log_{10} 188.81 = 45.52 \\text{dB}$\n",
    "\n",
    "Since this is a benchmark for machine learning models, we don't expect adversarial examples nor do we expect data drift. The data is as-is, nothing new would be added."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Q6: Trian your machine learner for accuracy at memory equivalent capacity\n",
    "\n",
    "We get MEC of 1038 bits in Q4, "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The decision layer has MEC 1056 bits\n",
      "Finished epoch 20\n",
      "Current LR: 0.001\n",
      "Total Loss 9.893238711811136e-06\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 782/782 [00:00<00:00, 3268.51it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Finished evaluating, memorized 49874 out of 50000\n",
      "\n",
      "Finished epoch 40\n",
      "Current LR: 0.001\n",
      "Total Loss 7.706622454861645e-06\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 782/782 [00:00<00:00, 3155.95it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Finished evaluating, memorized 49921 out of 50000\n",
      "\n",
      "Finished epoch 60\n",
      "Current LR: 0.001\n",
      "Total Loss 6.502566975541413e-06\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 782/782 [00:00<00:00, 3140.76it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Finished evaluating, memorized 49930 out of 50000\n",
      "\n",
      "Finished epoch 80\n",
      "Current LR: 0.001\n",
      "Total Loss 6.18038711763802e-06\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 782/782 [00:00<00:00, 3186.25it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Finished evaluating, memorized 49923 out of 50000\n",
      "\n",
      "Finished epoch 100\n",
      "Current LR: 0.001\n",
      "Total Loss 5.2952173064113595e-06\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 782/782 [00:00<00:00, 3194.10it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Finished evaluating, memorized 49944 out of 50000\n",
      "\n",
      "Finished epoch 120\n",
      "Current LR: 0.0001\n",
      "Total Loss 4.0207210076914635e-06\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 782/782 [00:00<00:00, 3076.17it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Finished evaluating, memorized 49950 out of 50000\n",
      "\n",
      "Finished epoch 140\n",
      "Current LR: 0.0001\n",
      "Total Loss 3.966811618738575e-06\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 782/782 [00:00<00:00, 3221.45it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Finished evaluating, memorized 49956 out of 50000\n",
      "\n",
      "Finished epoch 160\n",
      "Current LR: 0.0001\n",
      "Total Loss 3.901964191754814e-06\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 782/782 [00:00<00:00, 3174.16it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Finished evaluating, memorized 49949 out of 50000\n",
      "\n",
      "Finished epoch 180\n",
      "Current LR: 0.0001\n",
      "Total Loss 3.825977273663739e-06\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 782/782 [00:00<00:00, 3220.27it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Finished evaluating, memorized 49957 out of 50000\n",
      "\n",
      "Finished epoch 200\n",
      "Current LR: 0.0001\n",
      "Total Loss 3.7760869417979848e-06\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 782/782 [00:00<00:00, 3207.34it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Finished evaluating, memorized 49962 out of 50000\n",
      "\n",
      "Finished epoch 220\n",
      "Current LR: 1e-05\n",
      "Total Loss 3.6837961943092523e-06\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 782/782 [00:00<00:00, 3109.99it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Finished evaluating, memorized 49960 out of 50000\n",
      "\n",
      "Finished epoch 240\n",
      "Current LR: 1e-05\n",
      "Total Loss 3.6799933695874643e-06\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 782/782 [00:00<00:00, 3199.39it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Finished evaluating, memorized 49960 out of 50000\n",
      "\n",
      "Best trainning memorization is 0.99924\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "from linear_with_mec import create_linear_layer\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "net, mec = create_linear_layer(16)\n",
    "print('The decision layer has MEC {} bits'.format(mec))\n",
    "\n",
    "net.to(device)\n",
    "net.train()\n",
    "\n",
    "print_frequency = 50\n",
    "\n",
    "criterion = nn.CrossEntropyLoss().to(device)\n",
    "optimizer = torch.optim.Adam(net.parameters(), 1e-3)\n",
    "lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones=[100, 200])\n",
    "\n",
    "compressed_train_loader = DataLoader(CompressedCifar10(type='train'), 64, shuffle = True)\n",
    "\n",
    "def evaluate(model, data_loader):\n",
    "    model.eval()\n",
    "\n",
    "    correct_count = 0\n",
    "    total = 0\n",
    "    with torch.no_grad():\n",
    "        for data, label in tqdm(data_loader):\n",
    "            data = data.to(device)\n",
    "            label = label.to(device)\n",
    "\n",
    "            output = model(data)\n",
    "            _, predicted = output.max(1)\n",
    "            total += label.size(0)\n",
    "            correct_count += predicted.eq(label).sum().item()\n",
    "    return correct_count, total\n",
    "\n",
    "best_mem = 0\n",
    "\n",
    "def train(dataloader, model, optimizer, criterion):\n",
    "    model.train()\n",
    "    running_loss = 0\n",
    "    for data, label in dataloader:\n",
    "        data = data.to(device)\n",
    "        label = label.to(device)\n",
    "\n",
    "        output = model(data)\n",
    "        loss = criterion(output, label)\n",
    "        running_loss += loss\n",
    "\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "    return running_loss / len(dataloader)\n",
    "    \n",
    "for epoch in range(240):\n",
    "    running_loss =  train(compressed_train_loader, net, optimizer, criterion)\n",
    "\n",
    "    if epoch % 20 == 19:\n",
    "        print('Finished epoch {}'.format(epoch + 1))\n",
    "        print('Current LR: {}'.format(optimizer.param_groups[0]['lr']))\n",
    "        print('Total Loss {}'.format(running_loss / len(compressed_train_loader)))\n",
    "        correct_count, total = evaluate(net, compressed_train_loader)\n",
    "        print('Finished evaluating, memorized {} out of {}\\n'.format(correct_count, total))\n",
    "\n",
    "        best_mem = max(best_mem, correct_count / total)\n",
    "    lr_scheduler.step()\n",
    "\n",
    "print('Best trainning memorization is {}'.format(best_mem))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A decision layer with MEC 1056 bit (the data has 1038 bit) could memorize 99.9% of the trainning data in 240 epochs. (should investigate why we cannot get 100%???)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Q7: Trian for generalization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Testing for decision layer with MEC 3960 bits\n",
      "Finished trainning\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 157/157 [00:00<00:00, 3316.01it/s]\n",
      "100%|██████████| 782/782 [00:00<00:00, 3198.13it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Finished evaluating, predicted 9161 out of 10000\n",
      "\n",
      "Testing for decision layer with MEC 3630 bits\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "Cell \u001b[0;32mIn [70], line 16\u001b[0m\n\u001b[1;32m     13\u001b[0m lr_scheduler \u001b[39m=\u001b[39m torch\u001b[39m.\u001b[39moptim\u001b[39m.\u001b[39mlr_scheduler\u001b[39m.\u001b[39mMultiStepLR(optimizer, milestones\u001b[39m=\u001b[39m[\u001b[39m100\u001b[39m, \u001b[39m200\u001b[39m])\n\u001b[1;32m     15\u001b[0m \u001b[39mfor\u001b[39;00m epoch \u001b[39min\u001b[39;00m \u001b[39mrange\u001b[39m(\u001b[39m240\u001b[39m):\n\u001b[0;32m---> 16\u001b[0m     running_loss \u001b[39m=\u001b[39m  train(compressed_train_loader, net, optimizer, criterion)\n\u001b[1;32m     18\u001b[0m \u001b[39mprint\u001b[39m(\u001b[39m'\u001b[39m\u001b[39mFinished trainning\u001b[39m\u001b[39m'\u001b[39m)\n\u001b[1;32m     19\u001b[0m correct_test, total_test \u001b[39m=\u001b[39m evaluate(net, compressed_test_loader)\n",
      "Cell \u001b[0;32mIn [64], line 51\u001b[0m, in \u001b[0;36mtrain\u001b[0;34m(dataloader, model, optimizer, criterion)\u001b[0m\n\u001b[1;32m     48\u001b[0m     running_loss \u001b[39m+\u001b[39m\u001b[39m=\u001b[39m loss\n\u001b[1;32m     50\u001b[0m     optimizer\u001b[39m.\u001b[39mzero_grad()\n\u001b[0;32m---> 51\u001b[0m     loss\u001b[39m.\u001b[39;49mbackward()\n\u001b[1;32m     52\u001b[0m     optimizer\u001b[39m.\u001b[39mstep()\n\u001b[1;32m     53\u001b[0m \u001b[39mreturn\u001b[39;00m running_loss \u001b[39m/\u001b[39m \u001b[39mlen\u001b[39m(dataloader)\n",
      "File \u001b[0;32m/Library/Frameworks/Python.framework/Versions/3.10/lib/python3.10/site-packages/torch/_tensor.py:486\u001b[0m, in \u001b[0;36mTensor.backward\u001b[0;34m(self, gradient, retain_graph, create_graph, inputs)\u001b[0m\n\u001b[1;32m    476\u001b[0m \u001b[39mif\u001b[39;00m has_torch_function_unary(\u001b[39mself\u001b[39m):\n\u001b[1;32m    477\u001b[0m     \u001b[39mreturn\u001b[39;00m handle_torch_function(\n\u001b[1;32m    478\u001b[0m         Tensor\u001b[39m.\u001b[39mbackward,\n\u001b[1;32m    479\u001b[0m         (\u001b[39mself\u001b[39m,),\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    484\u001b[0m         inputs\u001b[39m=\u001b[39minputs,\n\u001b[1;32m    485\u001b[0m     )\n\u001b[0;32m--> 486\u001b[0m torch\u001b[39m.\u001b[39;49mautograd\u001b[39m.\u001b[39;49mbackward(\n\u001b[1;32m    487\u001b[0m     \u001b[39mself\u001b[39;49m, gradient, retain_graph, create_graph, inputs\u001b[39m=\u001b[39;49minputs\n\u001b[1;32m    488\u001b[0m )\n",
      "File \u001b[0;32m/Library/Frameworks/Python.framework/Versions/3.10/lib/python3.10/site-packages/torch/autograd/__init__.py:197\u001b[0m, in \u001b[0;36mbackward\u001b[0;34m(tensors, grad_tensors, retain_graph, create_graph, grad_variables, inputs)\u001b[0m\n\u001b[1;32m    192\u001b[0m     retain_graph \u001b[39m=\u001b[39m create_graph\n\u001b[1;32m    194\u001b[0m \u001b[39m# The reason we repeat same the comment below is that\u001b[39;00m\n\u001b[1;32m    195\u001b[0m \u001b[39m# some Python versions print out the first line of a multi-line function\u001b[39;00m\n\u001b[1;32m    196\u001b[0m \u001b[39m# calls in the traceback and some print out the last line\u001b[39;00m\n\u001b[0;32m--> 197\u001b[0m Variable\u001b[39m.\u001b[39;49m_execution_engine\u001b[39m.\u001b[39;49mrun_backward(  \u001b[39m# Calls into the C++ engine to run the backward pass\u001b[39;49;00m\n\u001b[1;32m    198\u001b[0m     tensors, grad_tensors_, retain_graph, create_graph, inputs,\n\u001b[1;32m    199\u001b[0m     allow_unreachable\u001b[39m=\u001b[39;49m\u001b[39mTrue\u001b[39;49;00m, accumulate_grad\u001b[39m=\u001b[39;49m\u001b[39mTrue\u001b[39;49;00m)\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "compressed_train_loader = DataLoader(CompressedCifar10(type='train'), 64, shuffle = True)\n",
    "compressed_test_loader = DataLoader(CompressedCifar10(type='test'), 64, shuffle = False)\n",
    "\n",
    "mec_acc = []\n",
    "\n",
    "for hidden_size in reversed(range(10, 65, 5)):\n",
    "    net, mec = create_linear_layer(hidden_size)\n",
    "    print('Testing for decision layer with MEC {} bits'.format(mec))\n",
    "\n",
    "\n",
    "    criterion = nn.CrossEntropyLoss().to(device)\n",
    "    optimizer = torch.optim.Adam(net.parameters(), 1e-3)\n",
    "    lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones=[100, 200])\n",
    "\n",
    "    for epoch in range(240):\n",
    "        running_loss =  train(compressed_train_loader, net, optimizer, criterion)\n",
    "\n",
    "    print('Finished trainning')\n",
    "    correct_test, total_test = evaluate(net, compressed_test_loader)\n",
    "    correct_train, total_train = evaluate(net, compressed_train_loader)\n",
    "    print('Finished evaluating, predicted {} out of {}\\n'.format(correct_test, total_test))\n",
    "\n",
    "    mec_acc.append((mec, correct_test / total_test, correct_train / total_train))\n",
    "\n",
    "print(mec_acc)\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Reference\n",
    "[1] https://arxiv.org/abs/2103.14749  \n",
    "[2] https://www.cs.toronto.edu/~kriz/learning-features-2009-TR.pdf"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.10.7 64-bit",
   "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.10.7"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "aee8b7b246df8f9039afb4144a1f6fd8d2ca17a180786b69acc140d282b71a49"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}