notebooks/Train-Merge-REPAIR-VGG11.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "9021242c",
   "metadata": {},
   "source": [
    "# Train-Merge-REPAIR-VGG11\n",
    "\n",
    "This is a minimal notebook which demonstrates REPAIR applied to a VGG11 network. It does the following:\n",
    "1. Separately trains two VGG11 networks on CIFAR-10, \"A\" and \"B\".\n",
    "2. Permutes the channels of each convolutional layer in \"B\" in order to align them with \"A\".\n",
    "3. Merges the two models in weight-space. The merged model performs poorly.\n",
    "4. Uses REPAIR to correct the neuronal statistics of the merged model.\n",
    "\n",
    "Notes:\n",
    "* The merged VGG network should initially attain 67.3% (+/-5%) accuracy. After REPAIR, it should reach 84.9% (+/-1%).\n",
    "* The trained networks should obtain 89-91% accuracy.\n",
    "* We use the original VGG architecture which does not contain normalization layers.\n",
    "* To align channels, we maximize correlations between the activations of matched neurons; this method is due to Li et al. (2015) https://arxiv.org/abs/1511.07543\n",
    "* REPAIR is a generalization of the method of resetting BatchNorms, which goes back to SWA https://arxiv.org/abs/1803.05407. We introduce it in https://arxiv.org/abs/2211.08403."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "c0db779d",
   "metadata": {},
   "outputs": [],
   "source": [
    "from tqdm import tqdm\n",
    "import numpy as np\n",
    "import scipy.optimize\n",
    "\n",
    "import torch\n",
    "from torch import nn\n",
    "from torch.cuda.amp import GradScaler, autocast\n",
    "from torch.nn import CrossEntropyLoss\n",
    "from torch.optim import SGD, lr_scheduler\n",
    "import torchvision\n",
    "import torchvision.transforms as T"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fdf1a55e",
   "metadata": {},
   "source": [
    "## Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "c477d1e9",
   "metadata": {},
   "outputs": [],
   "source": [
    "def save_model(model, i):\n",
    "    sd = model.state_dict()\n",
    "    torch.save(model.state_dict(), '%s.pt' % i)\n",
    "\n",
    "def load_model(model, i):\n",
    "    sd = torch.load('%s.pt' % i)\n",
    "    model.load_state_dict(sd)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "fabbc157",
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Files already downloaded and verified\n",
      "Files already downloaded and verified\n"
     ]
    }
   ],
   "source": [
    "CIFAR_MEAN = [125.307, 122.961, 113.8575]\n",
    "CIFAR_STD = [51.5865, 50.847, 51.255]\n",
    "normalize = T.Normalize(np.array(CIFAR_MEAN)/255, np.array(CIFAR_STD)/255)\n",
    "denormalize = T.Normalize(-np.array(CIFAR_MEAN)/np.array(CIFAR_STD), 255/np.array(CIFAR_STD))\n",
    "\n",
    "train_transform = T.Compose([\n",
    "    T.RandomHorizontalFlip(),\n",
    "    T.RandomCrop(32, padding=4),\n",
    "    T.ToTensor(),\n",
    "    normalize,\n",
    "])\n",
    "test_transform = T.Compose([\n",
    "    T.ToTensor(),\n",
    "    normalize,\n",
    "])\n",
    "train_dset = torchvision.datasets.CIFAR10(root='/tmp', train=True,\n",
    "                                        download=True, transform=train_transform)\n",
    "test_dset = torchvision.datasets.CIFAR10(root='/tmp', train=False,\n",
    "                                        download=True, transform=test_transform)\n",
    "\n",
    "train_aug_loader = torch.utils.data.DataLoader(train_dset, batch_size=500, shuffle=True, num_workers=8)\n",
    "test_loader = torch.utils.data.DataLoader(test_dset, batch_size=500, shuffle=False, num_workers=8)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "31e22780",
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluate(model):\n",
    "    model.eval()\n",
    "    correct = 0\n",
    "    with torch.no_grad(), autocast():\n",
    "        for inputs, labels in test_loader:\n",
    "            outputs = model(inputs.cuda())\n",
    "            pred = outputs.argmax(dim=1)\n",
    "            correct += (labels.cuda() == pred).sum().item()\n",
    "    return correct / len(test_loader.dataset)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e42cdbf3",
   "metadata": {},
   "source": [
    "# Train two VGG11 networks on CIFAR-10"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "51b6c738",
   "metadata": {},
   "outputs": [],
   "source": [
    "# https://github.com/kuangliu/pytorch-cifar/blob/master/models/vgg.py\n",
    "cfg = {\n",
    "    'VGG11': [64, 'M', 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'],\n",
    "    'VGG13': [64, 64, 'M', 128, 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'],\n",
    "    'VGG16': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'M', 512, 512, 512, 'M', 512, 512, 512, 'M'],\n",
    "    'VGG19': [64, 64, 'M', 128, 128, 'M', 256, 256, 256, 256, 'M', 512, 512, 512, 512, 'M', 512, 512, 512, 512, 'M'],\n",
    "}\n",
    "\n",
    "class VGG(nn.Module):\n",
    "    def __init__(self, vgg_name, w=1):\n",
    "        super(VGG, self).__init__()\n",
    "        self.w = w\n",
    "        self.features = self._make_layers(cfg[vgg_name])\n",
    "        self.classifier = nn.Linear(self.w*512, 10)\n",
    "\n",
    "    def forward(self, x):\n",
    "        out = self.features(x)\n",
    "        out = out.view(out.size(0), -1)\n",
    "        out = self.classifier(out)\n",
    "        return out\n",
    "\n",
    "    def _make_layers(self, cfg):\n",
    "        layers = []\n",
    "        in_channels = 3\n",
    "        for x in cfg:\n",
    "            if x == 'M':\n",
    "                layers += [nn.MaxPool2d(kernel_size=2, stride=2)]\n",
    "            else:\n",
    "                layers.append(nn.Conv2d(in_channels if in_channels == 3 else self.w*in_channels,\n",
    "                                     self.w*x, kernel_size=3, padding=1))\n",
    "                layers.append(nn.ReLU(inplace=True))\n",
    "                in_channels = x\n",
    "        layers += [nn.AvgPool2d(kernel_size=1, stride=1)]\n",
    "        return nn.Sequential(*layers)\n",
    "    \n",
    "def vgg11(w=1):\n",
    "    return VGG('VGG11', w).cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "d120e6b8",
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_model(w=1):\n",
    "    model = vgg11(w)\n",
    "    optimizer = SGD(model.parameters(), lr=0.08, momentum=0.9, weight_decay=5e-4)\n",
    "\n",
    "    EPOCHS = 100\n",
    "    ne_iters = len(train_aug_loader)\n",
    "    lr_schedule = np.interp(np.arange(1+EPOCHS*ne_iters), [0, 5*ne_iters, EPOCHS*ne_iters], [0, 1, 0])\n",
    "    scheduler = lr_scheduler.LambdaLR(optimizer, lr_schedule.__getitem__)\n",
    "\n",
    "    scaler = GradScaler()\n",
    "    loss_fn = CrossEntropyLoss()\n",
    "\n",
    "    losses = []\n",
    "    for epoch in tqdm(range(EPOCHS)):\n",
    "        model.train()\n",
    "        for i, (inputs, labels) in enumerate(train_aug_loader):\n",
    "            optimizer.zero_grad(set_to_none=True)\n",
    "            with autocast():\n",
    "                outputs = model(inputs.cuda())\n",
    "                loss = loss_fn(outputs, labels.cuda())\n",
    "            scaler.scale(loss).backward()\n",
    "            scaler.step(optimizer)\n",
    "            scaler.update()\n",
    "            scheduler.step()\n",
    "            losses.append(loss.item())\n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "08a4ae3c",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 100/100 [03:38<00:00,  2.19s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.8937\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 100/100 [04:13<00:00,  2.53s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.8981\n"
     ]
    }
   ],
   "source": [
    "model = train_model()\n",
    "print(evaluate(model))\n",
    "save_model(model, 'vgg11_v1')\n",
    "\n",
    "model = train_model()\n",
    "print(evaluate(model))\n",
    "save_model(model, 'vgg11_v2')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d7114582",
   "metadata": {},
   "source": [
    "# Permute the channels of model B to align with model A"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "id": "363e68ea",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Given two networks net0, net1 which each output a feature map of shape NxCxWxH,\n",
    "# this will reshape both outputs to (N*W*H)xC\n",
    "# and then compute a CxC correlation matrix between the two\n",
    "def run_corr_matrix(net0, net1):\n",
    "    n = len(train_aug_loader)\n",
    "    with torch.no_grad():\n",
    "        net0.eval()\n",
    "        net1.eval()\n",
    "        for i, (images, _) in enumerate(tqdm(train_aug_loader)):\n",
    "            \n",
    "            img_t = images.float().cuda()\n",
    "            out0 = net0(img_t).double()\n",
    "            out0 = out0.permute(0, 2, 3, 1).reshape(-1, out0.shape[1])\n",
    "            out1 = net1(img_t).double()\n",
    "            out1 = out1.permute(0, 2, 3, 1).reshape(-1, out1.shape[1])\n",
    "\n",
    "            # save batchwise first+second moments and outer product\n",
    "            mean0_b = out0.mean(dim=0)\n",
    "            mean1_b = out1.mean(dim=0)\n",
    "            sqmean0_b = out0.square().mean(dim=0)\n",
    "            sqmean1_b = out1.square().mean(dim=0)\n",
    "            outer_b = (out0.T @ out1) / out0.shape[0]\n",
    "            if i == 0:\n",
    "                mean0 = torch.zeros_like(mean0_b)\n",
    "                mean1 = torch.zeros_like(mean1_b)\n",
    "                sqmean0 = torch.zeros_like(sqmean0_b)\n",
    "                sqmean1 = torch.zeros_like(sqmean1_b)\n",
    "                outer = torch.zeros_like(outer_b)\n",
    "            mean0 += mean0_b / n\n",
    "            mean1 += mean1_b / n\n",
    "            sqmean0 += sqmean0_b / n\n",
    "            sqmean1 += sqmean1_b / n\n",
    "            outer += outer_b / n\n",
    "\n",
    "    cov = outer - torch.outer(mean0, mean1)\n",
    "    std0 = (sqmean0 - mean0**2).sqrt()\n",
    "    std1 = (sqmean1 - mean1**2).sqrt()\n",
    "    corr = cov / (torch.outer(std0, std1) + 1e-4)\n",
    "    return corr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "id": "1f9ef263",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_layer_perm1(corr_mtx):\n",
    "    corr_mtx_a = corr_mtx.cpu().numpy()\n",
    "    row_ind, col_ind = scipy.optimize.linear_sum_assignment(corr_mtx_a, maximize=True)\n",
    "    assert (row_ind == np.arange(len(corr_mtx_a))).all()\n",
    "    perm_map = torch.tensor(col_ind).long()\n",
    "    return perm_map\n",
    "\n",
    "# returns the channel-permutation to make layer1's activations most closely\n",
    "# match layer0's.\n",
    "def get_layer_perm(net0, net1):\n",
    "    corr_mtx = run_corr_matrix(net0, net1)\n",
    "    return get_layer_perm1(corr_mtx)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "id": "f1aae9be",
   "metadata": {},
   "outputs": [],
   "source": [
    "# modifies the weight matrices of a convolution and batchnorm\n",
    "# layer given a permutation of the output channels\n",
    "def permute_output(perm_map, layer):\n",
    "    pre_weights = [layer.weight,\n",
    "                   layer.bias]\n",
    "    for w in pre_weights:\n",
    "        w.data = w[perm_map]\n",
    "\n",
    "# modifies the weight matrix of a layer for a given permutation of the input channels\n",
    "# works for both conv2d and linear\n",
    "def permute_input(perm_map, layer):\n",
    "    w = layer.weight\n",
    "    w.data = w[:, perm_map]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "id": "c0db278e",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(0.8937, 0.8981)"
      ]
     },
     "execution_count": 47,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model0 = vgg11()\n",
    "model1 = vgg11()\n",
    "load_model(model0, 'vgg11_v1')\n",
    "load_model(model1, 'vgg11_v2')\n",
    "\n",
    "evaluate(model0), evaluate(model1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "id": "b3f2f3bb",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 100/100 [00:02<00:00, 41.18it/s]\n",
      "100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 100/100 [00:02<00:00, 40.86it/s]\n",
      "100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 100/100 [00:02<00:00, 44.91it/s]\n",
      "100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 100/100 [00:02<00:00, 46.51it/s]\n",
      "100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 100/100 [00:02<00:00, 42.78it/s]\n",
      "100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 100/100 [00:02<00:00, 41.18it/s]\n",
      "100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 100/100 [00:03<00:00, 31.86it/s]\n",
      "100%|███████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 100/100 [00:02<00:00, 34.22it/s]\n"
     ]
    }
   ],
   "source": [
    "def subnet(model, n_layers):\n",
    "    return model.features[:n_layers]\n",
    "\n",
    "feats1 = model1.features\n",
    "\n",
    "n = len(feats1)\n",
    "for i in range(n):\n",
    "    if not isinstance(feats1[i], nn.Conv2d):\n",
    "        continue\n",
    "    \n",
    "    # permute the outputs of the current conv layer\n",
    "    assert isinstance(feats1[i+1], nn.ReLU)\n",
    "    perm_map = get_layer_perm(subnet(model0, i+2), subnet(model1, i+2))\n",
    "    permute_output(perm_map, feats1[i])\n",
    "    \n",
    "    # look for the next conv layer, whose inputs should be permuted the same way\n",
    "    next_layer = None\n",
    "    for j in range(i+1, n):\n",
    "        if isinstance(feats1[j], nn.Conv2d):\n",
    "            next_layer = feats1[j]\n",
    "            break\n",
    "    if next_layer is None:\n",
    "        next_layer = model1.classifier\n",
    "    permute_input(perm_map, next_layer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "id": "3112ee07",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.8981\n"
     ]
    }
   ],
   "source": [
    "# ensure accuracy didn't change\n",
    "# (it may be slightly different due to non-associativity of floating point arithmetic)\n",
    "print(evaluate(model1))\n",
    "save_model(model1, 'vgg11_v2_perm1')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bc494077",
   "metadata": {},
   "source": [
    "# Merge the two networks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "id": "1a7ae6a7",
   "metadata": {},
   "outputs": [],
   "source": [
    "def mix_weights(net, alpha, key0, key1):\n",
    "    sd0 = torch.load('%s.pt' % key0)\n",
    "    sd1 = torch.load('%s.pt' % key1)\n",
    "    sd_alpha = {k: (1 - alpha) * sd0[k].cuda() + alpha * sd1[k].cuda()\n",
    "                for k in sd0.keys()}\n",
    "    net.load_state_dict(sd_alpha)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "id": "00b6b3e5",
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(α=0): 89.4% \t\t<-- Model A\n",
      "(α=1): 89.8% \t\t<-- Model B\n",
      "(α=0.5): 67.3% \t\t<-- Merged model\n"
     ]
    }
   ],
   "source": [
    "k0 = 'vgg11_v1'\n",
    "k1 = 'vgg11_v2_perm1'\n",
    "model0 = vgg11()\n",
    "model1 = vgg11()\n",
    "model_a = vgg11()\n",
    "mix_weights(model0, 0.0, k0, k1)\n",
    "mix_weights(model1, 1.0, k0, k1)\n",
    "\n",
    "alpha = 0.5\n",
    "mix_weights(model_a, alpha, k0, k1)\n",
    "print('(α=0): %.1f%% \\t\\t<-- Model A' % (100*evaluate(model0)))\n",
    "print('(α=1): %.1f%% \\t\\t<-- Model B' % (100*evaluate(model1)))\n",
    "print('(α=0.5): %.1f%% \\t\\t<-- Merged model' % (100*evaluate(model_a)))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "89396da3",
   "metadata": {},
   "source": [
    "# Correct neuronal statistics with REPAIR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "id": "4ff057c6",
   "metadata": {},
   "outputs": [],
   "source": [
    "class TrackLayer(nn.Module):\n",
    "    def __init__(self, layer):\n",
    "        super().__init__()\n",
    "        self.layer = layer\n",
    "        self.bn = nn.BatchNorm2d(layer.out_channels)\n",
    "        \n",
    "    def get_stats(self):\n",
    "        return (self.bn.running_mean, self.bn.running_var.sqrt())\n",
    "        \n",
    "    def forward(self, x):\n",
    "        x1 = self.layer(x)\n",
    "        self.bn(x1)\n",
    "        return x1\n",
    "\n",
    "class ResetLayer(nn.Module):\n",
    "    def __init__(self, layer):\n",
    "        super().__init__()\n",
    "        self.layer = layer\n",
    "        self.bn = nn.BatchNorm2d(layer.out_channels)\n",
    "        \n",
    "    def set_stats(self, goal_mean, goal_std):\n",
    "        self.bn.bias.data = goal_mean\n",
    "        self.bn.weight.data = goal_std\n",
    "        \n",
    "    def forward(self, x):\n",
    "        x1 = self.layer(x)\n",
    "        return self.bn(x1)\n",
    "\n",
    "# adds TrackLayers around every conv layer\n",
    "def make_tracked_net(net):\n",
    "    net1 = vgg11()\n",
    "    net1.load_state_dict(net.state_dict())\n",
    "    for i, layer in enumerate(net1.features):\n",
    "        if isinstance(layer, nn.Conv2d):\n",
    "            net1.features[i] = TrackLayer(layer)\n",
    "    return net1.cuda().eval()\n",
    "\n",
    "# adds ResetLayers around every conv layer\n",
    "def make_repaired_net(net):\n",
    "    net1 = vgg11()\n",
    "    net1.load_state_dict(net.state_dict())\n",
    "    for i, layer in enumerate(net1.features):\n",
    "        if isinstance(layer, nn.Conv2d):\n",
    "            net1.features[i] = ResetLayer(layer)\n",
    "    return net1.cuda().eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "id": "c5dce4e3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# reset all tracked BN stats against training data\n",
    "def reset_bn_stats(model):\n",
    "    # resetting stats to baseline first as below is necessary for stability\n",
    "    for m in model.modules():\n",
    "        if type(m) == nn.BatchNorm2d:\n",
    "            m.momentum = None # use simple average\n",
    "            m.reset_running_stats()\n",
    "    model.train()\n",
    "    with torch.no_grad(), autocast():\n",
    "        for images, _ in train_aug_loader:\n",
    "            output = model(images.cuda())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "id": "f4fce3fe",
   "metadata": {},
   "outputs": [],
   "source": [
    "## Calculate all neuronal statistics in the endpoint networks\n",
    "wrap0 = make_tracked_net(model0)\n",
    "wrap1 = make_tracked_net(model1)\n",
    "reset_bn_stats(wrap0)\n",
    "reset_bn_stats(wrap1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "id": "d467bdee",
   "metadata": {},
   "outputs": [],
   "source": [
    "wrap_a = make_repaired_net(model_a)\n",
    "# Iterate through corresponding triples of (TrackLayer, TrackLayer, ResetLayer)\n",
    "# around conv layers in (model0, model1, model_a).\n",
    "for track0, track1, reset_a in zip(wrap0.modules(), wrap1.modules(), wrap_a.modules()): \n",
    "    if not isinstance(track0, TrackLayer):\n",
    "        continue  \n",
    "    assert (isinstance(track0, TrackLayer)\n",
    "            and isinstance(track1, TrackLayer)\n",
    "            and isinstance(reset_a, ResetLayer))\n",
    "\n",
    "    # get neuronal statistics of original networks\n",
    "    mu0, std0 = track0.get_stats()\n",
    "    mu1, std1 = track1.get_stats()\n",
    "    # set the goal neuronal statistics for the merged network \n",
    "    goal_mean = (1 - alpha) * mu0 + alpha * mu1\n",
    "    goal_std = (1 - alpha) * std0 + alpha * std1\n",
    "    reset_a.set_stats(goal_mean, goal_std)\n",
    "\n",
    "# Estimate mean/vars such that when added BNs are set to eval mode,\n",
    "# neuronal stats will be goal_mean and goal_std.\n",
    "reset_bn_stats(wrap_a)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3313ef52",
   "metadata": {},
   "source": [
    "### Fuse added BNs so that weights are compatible with original architecture"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "id": "693446dc",
   "metadata": {},
   "outputs": [],
   "source": [
    "def fuse_conv_bn(conv, bn):\n",
    "    fused_conv = torch.nn.Conv2d(conv.in_channels,\n",
    "                                 conv.out_channels,\n",
    "                                 kernel_size=conv.kernel_size,\n",
    "                                 stride=conv.stride,\n",
    "                                 padding=conv.padding,\n",
    "                                 bias=True)\n",
    "\n",
    "    # set weights\n",
    "    w_conv = conv.weight.clone()\n",
    "    bn_std = (bn.eps + bn.running_var).sqrt()\n",
    "    gamma = bn.weight / bn_std\n",
    "    fused_conv.weight.data = (w_conv * gamma.reshape(-1, 1, 1, 1))\n",
    "\n",
    "    # set bias\n",
    "    beta = bn.bias + gamma * (-bn.running_mean + conv.bias)\n",
    "    fused_conv.bias.data = beta\n",
    "    \n",
    "    return fused_conv\n",
    "\n",
    "def fuse_tracked_net(net):\n",
    "    net1 = vgg11()\n",
    "    for i, rlayer in enumerate(net.features):\n",
    "        if isinstance(rlayer, ResetLayer):\n",
    "            fused_conv = fuse_conv_bn(rlayer.layer, rlayer.bn)\n",
    "            net1.features[i].load_state_dict(fused_conv.state_dict())\n",
    "    net1.classifier.load_state_dict(net.classifier.state_dict())\n",
    "    return net1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "id": "d5d81c46",
   "metadata": {},
   "outputs": [],
   "source": [
    "# fuse the rescaling+shift coefficients back into conv layers\n",
    "model_b = fuse_tracked_net(wrap_a)\n",
    "assert abs(evaluate(model_b) - evaluate(wrap_a)) < 0.005"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "id": "2d899359",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(α=0): 89.4% \t\t<-- Model A\n",
      "(α=1): 89.8% \t\t<-- Model B\n",
      "(α=0.5): 84.9% \t\t<-- Merged model with REPAIR\n"
     ]
    }
   ],
   "source": [
    "print('(α=0): %.1f%% \\t\\t<-- Model A' % (100*evaluate(model0)))\n",
    "print('(α=1): %.1f%% \\t\\t<-- Model B' % (100*evaluate(model1)))\n",
    "print('(α=0.5): %.1f%% \\t\\t<-- Merged model with REPAIR' % (100*evaluate(model_b)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "347efbc5",
   "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.9.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}