hw7_bn_remove.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 0
    },
    "id": "SpSFfApAMFE6",
    "outputId": "5a9d9f33-f339-4e74-ec11-8672ad529827"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Files already downloaded and verified\n",
      "tensor(7.8475e-09)\n"
     ]
    }
   ],
   "source": [
    "import torch.nn as nn\n",
    "from torch.utils.data import DataLoader\n",
    "import torch\n",
    "import torchvision\n",
    "import numpy as np\n",
    "import torchvision.transforms as transforms\n",
    "\n",
    "\n",
    "# Instantiate model with BN and load trained parameters\n",
    "class smallNetTrain(nn.Module) :\n",
    "    # CIFAR-10 data is 32*32 images with 3 RGB channels\n",
    "    def __init__(self, input_dim=3*32*32) :\n",
    "        super().__init__()\n",
    "\n",
    "        self.conv1 = nn.Sequential(\n",
    "                            nn.Conv2d(3, 16, kernel_size=3, padding=1),\n",
    "                            nn.BatchNorm2d(16),\n",
    "                            nn.ReLU()\n",
    "                            )      \n",
    "        self.conv2 = nn.Sequential(\n",
    "                            nn.Conv2d(16, 16, kernel_size=3, padding=1),\n",
    "                            nn.BatchNorm2d(16),\n",
    "                            nn.ReLU()\n",
    "                            ) \n",
    "        self.fc1 = nn.Sequential(\n",
    "                            nn.Linear(16*32*32, 32*32),\n",
    "                            nn.BatchNorm1d(32*32),\n",
    "                            nn.ReLU()\n",
    "                            )   \n",
    "        self.fc2 = nn.Sequential(\n",
    "                            nn.Linear(32*32, 10),\n",
    "                            nn.ReLU()\n",
    "                            )   \n",
    "    def forward(self, x) :\n",
    "        x = self.conv1(x)\n",
    "        x = self.conv2(x)\n",
    "        x = x.float().view(-1, 16*32*32)\n",
    "        x = self.fc1(x)\n",
    "        x = self.fc2(x)\n",
    "        \n",
    "        return x\n",
    "    \n",
    "model = smallNetTrain()\n",
    "model.load_state_dict(torch.load(\"./smallNetSaved\",map_location=torch.device('cpu')))\n",
    "\n",
    "\n",
    "# Instantiate model without BN\n",
    "class smallNetTest(nn.Module) :\n",
    "    # CIFAR-10 data is 32*32 images with 3 RGB channels\n",
    "    def __init__(self, input_dim=3*32*32) :\n",
    "        super().__init__()\n",
    "\n",
    "        self.conv1 = nn.Sequential(\n",
    "                            nn.Conv2d(3, 16, kernel_size=3, padding=1),\n",
    "                            nn.ReLU()\n",
    "                            )      \n",
    "        self.conv2 = nn.Sequential(\n",
    "                            nn.Conv2d(16, 16, kernel_size=3, padding=1),\n",
    "                            nn.ReLU()\n",
    "                            ) \n",
    "        self.fc1 = nn.Sequential(\n",
    "                            nn.Linear(16*32*32, 32*32),\n",
    "                            nn.ReLU()\n",
    "                            )   \n",
    "        self.fc2 = nn.Sequential(\n",
    "                            nn.Linear(32*32, 10),\n",
    "                            nn.ReLU()\n",
    "                            )   \n",
    "    def forward(self, x) :\n",
    "        x = self.conv1(x)\n",
    "        x = self.conv2(x)\n",
    "        x = x.float().view(-1, 16*32*32)\n",
    "        x = self.fc1(x)\n",
    "        x = self.fc2(x)\n",
    "        \n",
    "        return x\n",
    "    \n",
    "model_test = smallNetTest()\n",
    "\n",
    "\n",
    "\n",
    "# Initialize weights of model without BN\n",
    "\n",
    "conv1_bn_beta, conv1_bn_gamma = model.conv1[1].bias, model.conv1[1].weight\n",
    "conv1_bn_mean, conv1_bn_var = model.conv1[1].running_mean, model.conv1[1].running_var\n",
    "conv2_bn_beta, conv2_bn_gamma = model.conv2[1].bias, model.conv2[1].weight\n",
    "conv2_bn_mean, conv2_bn_var = model.conv2[1].running_mean, model.conv2[1].running_var\n",
    "fc1_bn_beta, fc1_bn_gamma = model.fc1[1].bias, model.fc1[1].weight\n",
    "fc1_bn_mean, fc1_bn_var = model.fc1[1].running_mean, model.fc1[1].running_var\n",
    "eps = 1e-05\n",
    "\n",
    "'''\n",
    "print(conv1_bn_beta.shape, conv1_bn_gamma.shape) # 16, 16\n",
    "print(conv2_bn_beta.shape, conv2_bn_gamma.shape) # 16, 16\n",
    "print(fc1_bn_beta.shape, fc1_bn_gamma.shape) # 1024, 1024\n",
    "print(model.conv1[0].weight.shape) # 16,3,3,3\n",
    "'''\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "# Initialize the following parameters\n",
    "model_test.conv1[0].weight.data = conv1_bn_gamma.unsqueeze(-1).unsqueeze(-1).unsqueeze(-1) * model.conv1[0].weight / np.sqrt(conv1_bn_var.unsqueeze(-1).unsqueeze(-1).unsqueeze(-1)+eps)\n",
    "\n",
    "model_test.conv1[0].bias.data = conv1_bn_beta + conv1_bn_gamma * (model.conv1[0].bias - conv1_bn_mean)/np.sqrt(conv1_bn_var+eps)\n",
    "\n",
    "model_test.conv2[0].weight.data = conv2_bn_gamma.unsqueeze(-1).unsqueeze(-1).unsqueeze(-1) * model.conv2[0].weight / np.sqrt(conv2_bn_var.unsqueeze(-1).unsqueeze(-1).unsqueeze(-1)+eps)\n",
    "\n",
    "model_test.conv2[0].bias.data = conv2_bn_beta + conv2_bn_gamma * (model.conv2[0].bias - conv2_bn_mean)/np.sqrt(conv2_bn_var+eps)\n",
    "\n",
    "model_test.fc1[0].weight.data = fc1_bn_gamma.unsqueeze(-1) * model.fc1[0].weight / np.sqrt(fc1_bn_var.unsqueeze(-1)+eps)\n",
    "\n",
    "model_test.fc1[0].bias.data = fc1_bn_beta + fc1_bn_gamma * (model.fc1[0].bias - fc1_bn_mean)/np.sqrt(fc1_bn_var+eps)\n",
    "\n",
    "model_test.fc2[0].weight.data = model.fc2[0].weight.data\n",
    "\n",
    "model_test.fc2[0].bias.data = model.fc2[0].bias.data\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "# Verify difference between model and model_test\n",
    "\n",
    "model.eval()  \n",
    "# model_test.eval()  # not necessary since model_test has no BN or dropout \n",
    "\n",
    "\n",
    "test_dataset = torchvision.datasets.CIFAR10(root='./cifar_10data/',\n",
    "                                train=False, \n",
    "                                transform=transforms.ToTensor(), download = True)\n",
    "test_loader = torch.utils.data.DataLoader(dataset=test_dataset, batch_size=100, shuffle=False)\n",
    "\n",
    "diff = []\n",
    "with torch.no_grad():\n",
    "    for images, _ in test_loader:\n",
    "        diff.append(torch.norm(model(images) - model_test(images))**2) \n",
    "        \n",
    "print(max(diff)) # If less than 1e-08, you got the right answer.\n",
    "\n"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "provenance": []
  },
  "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.9.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}