importance_sampling_SVRG_k.ipynb

{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "nOQafPJJfLRy",
        "outputId": "902c49cc-8894-4dd0-d878-0a7d41e76197"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Collecting utils\n",
            "  Downloading utils-1.0.2.tar.gz (13 kB)\n",
            "  Preparing metadata (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "Building wheels for collected packages: utils\n",
            "  Building wheel for utils (setup.py) ... \u001b[?25l\u001b[?25hdone\n",
            "  Created wheel for utils: filename=utils-1.0.2-py2.py3-none-any.whl size=13905 sha256=cc6ffe52789905371da71be531d759bb81f9a3cdce1318f8abe05c744ebb753c\n",
            "  Stored in directory: /root/.cache/pip/wheels/b8/39/f5/9d0ca31dba85773ececf0a7f5469f18810e1c8a8ed9da28ca7\n",
            "Successfully built utils\n",
            "Installing collected packages: utils\n",
            "Successfully installed utils-1.0.2\n"
          ]
        }
      ],
      "source": [
        "!pip install utils"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "v6mV9NL-fE54"
      },
      "outputs": [],
      "source": [
        "from torch.utils.data import DataLoader\n",
        "import numpy as np\n",
        "import matplotlib.pyplot as plt\n",
        "import torchvision\n",
        "from tqdm import tqdm\n",
        "import utils\n",
        "import torch\n",
        "from torchvision import transforms\n",
        "import torch.nn as nn\n",
        "import torch.nn.functional as F"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5scJYHifURI8"
      },
      "outputs": [],
      "source": [
        "class CNN(nn.Module):\n",
        "    def __init__(self, num_classes=10):\n",
        "        super(CNN, self).__init__()\n",
        "        self.features = nn.Sequential(\n",
        "            # BEGIN Solution (do not delete this comment!)\n",
        "\n",
        "            nn.Conv2d(3, 64, kernel_size=3, padding=1),\n",
        "            nn.BatchNorm2d(64),\n",
        "            nn.Conv2d(64, 64, kernel_size=5, padding=2),\n",
        "            nn.ReLU(),\n",
        "            nn.MaxPool2d(kernel_size=2, stride=2),\n",
        "            nn.BatchNorm2d(64),\n",
        "\n",
        "            nn.Conv2d(64, 128, kernel_size=3, padding=1),\n",
        "            nn.BatchNorm2d(128),\n",
        "            nn.Conv2d(128, 128, kernel_size=5, padding=2),\n",
        "            nn.ReLU(),\n",
        "            nn.MaxPool2d(kernel_size=2, stride=2),\n",
        "            nn.BatchNorm2d(128),\n",
        "\n",
        "            nn.Conv2d(128, 256, kernel_size=3, padding=1),\n",
        "            nn.BatchNorm2d(256),\n",
        "            nn.ReLU(),\n",
        "            nn.MaxPool2d(kernel_size=2, stride=2),\n",
        "\n",
        "            # END Solution (do not delete this comment!)\n",
        "        )\n",
        "        self.classifier = nn.Sequential(\n",
        "            # BEGIN Solution (do not delete this comment!)\n",
        "\n",
        "            nn.Linear(4096, 512),\n",
        "            nn.BatchNorm1d(512),\n",
        "            nn.ReLU(),\n",
        "            nn.Dropout(0.2),\n",
        "\n",
        "            nn.Linear(512, 256),\n",
        "            nn.BatchNorm1d(256),\n",
        "            nn.ReLU(),\n",
        "            nn.Dropout(0.2),\n",
        "\n",
        "            nn.Linear(256, num_classes)\n",
        "\n",
        "            # END Solution (do not delete this comment!)\n",
        "        )\n",
        "\n",
        "    def forward(self, x):\n",
        "        # BEGIN Solution (do not delete this comment!)\n",
        "\n",
        "        x = self.features(x)\n",
        "        x = x.view(x.size(0), -1)\n",
        "        x = self.classifier(x)\n",
        "        return x"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jsQS2RXt58rl"
      },
      "outputs": [],
      "source": [
        "def epoch_train(loader, clf, criterion, opt):\n",
        "\n",
        "    device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')\n",
        "    clf = clf.to(device)\n",
        "\n",
        "    clf.train(True)\n",
        "    total_loss = 0\n",
        "    correct = 0\n",
        "    total = 0\n",
        "    for model_input, target in loader:\n",
        "\n",
        "        model_input = model_input.to(device)\n",
        "        target = target.to(device)\n",
        "\n",
        "        model_output = clf(model_input)\n",
        "        loss = criterion(model_output, target)\n",
        "\n",
        "        #accuracy = utils.get_accuracy(clf, loader, device)\n",
        "\n",
        "        opt.zero_grad()\n",
        "        loss.backward()\n",
        "        opt.step()\n",
        "\n",
        "        total_loss += loss.item()\n",
        "        _, predicted = model_output.max(1)\n",
        "        total += target.size(0)\n",
        "        correct += predicted.eq(target).sum().item()\n",
        "\n",
        "    avg_loss = total_loss / len(loader)\n",
        "    avg_accuracy = correct / total\n",
        "    return avg_loss, avg_accuracy\n",
        "\n",
        "def epoch_test(loader, clf, criterion):\n",
        "\n",
        "    clf.eval()\n",
        "    device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')\n",
        "    clf = clf.to(device)\n",
        "    total_loss = 0\n",
        "    correct = 0\n",
        "    total = 0\n",
        "    for model_input, target in loader:\n",
        "        model_input = model_input.to(device)\n",
        "        target = target.to(device)\n",
        "\n",
        "        model_output = clf(model_input)\n",
        "        loss = criterion(model_output, target)\n",
        "\n",
        "        total_loss += loss.item()\n",
        "        _, predicted = model_output.max(1)\n",
        "        total += target.size(0)\n",
        "        correct += predicted.eq(target).sum().item()\n",
        "\n",
        "    avg_loss = total_loss / len(loader)\n",
        "    avg_accuracy = correct / total\n",
        "    return avg_loss, avg_accuracy\n",
        "\n",
        "def train(train_loader, test_loader, clf, criterion, opt, n_epochs=50):\n",
        "    for epoch in tqdm(range(n_epochs)):\n",
        "        train_loss, train_acc = epoch_train(train_loader, clf, criterion, opt)\n",
        "        test_loss, test_acc = epoch_test(test_loader, clf, criterion)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ModWsWJrAq91"
      },
      "source": [
        "# WideResNet Translate"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "sqV-B9MjfF-O"
      },
      "outputs": [],
      "source": [
        "class WideResNet(nn.Module):\n",
        "    def __init__(self, depth, widen_factor, input_shape, output_size, drop_rate=0.0):\n",
        "        super(WideResNet, self).__init__()\n",
        "\n",
        "        def conv(channels, strides):\n",
        "            return nn.Sequential(\n",
        "                nn.BatchNorm2d(channels),\n",
        "                nn.ReLU(),\n",
        "                nn.Conv2d(channels, channels, kernel_size=3, stride=strides, padding=1, bias=False),\n",
        "                nn.Dropout(drop_rate) if drop_rate > 0 else nn.Identity(),\n",
        "                nn.BatchNorm2d(channels),\n",
        "                nn.ReLU(),\n",
        "                nn.Conv2d(channels, channels, kernel_size=3, stride=1, padding=1, bias=False)\n",
        "            )\n",
        "\n",
        "        def resize(x, shape):\n",
        "            if x.size() == shape:\n",
        "                return x\n",
        "            channels = shape[1]\n",
        "            strides = x.size(2) // shape[2]\n",
        "            return nn.Conv2d(x.size(1), channels, kernel_size=1, stride=strides, padding=0, bias=False)(x)\n",
        "\n",
        "        def block(channels, widen_factor, n, strides):\n",
        "            layers = []\n",
        "            for i in range(n):\n",
        "                layers.append(conv(channels*widen_factor, strides if i == 0 else 1))\n",
        "                layers.append(nn.ReLU())\n",
        "            return nn.Sequential(*layers)\n",
        "\n",
        "        n = int((depth-4)/6)\n",
        "\n",
        "        self.group0 = nn.Conv2d(input_shape[0], 16, kernel_size=3, stride=1, padding=1, bias=False)\n",
        "        self.group1 = block(16, widen_factor, n, 1)\n",
        "        self.group2 = block(32, widen_factor, n, 2)\n",
        "        self.group3 = block(64, widen_factor, n, 2)\n",
        "\n",
        "        self.pool = nn.AdaptiveAvgPool2d((1, 1))\n",
        "        self.fc = nn.Linear(64, output_size)\n",
        "\n",
        "    def forward(self, x):\n",
        "        x = self.group0(x)\n",
        "        x = self.group1(x)\n",
        "        x = self.group2(x)\n",
        "        x = self.group3(x)\n",
        "\n",
        "        x = F.relu(x)\n",
        "        x = self.pool(x)\n",
        "        x = x.view(x.size(0), -1)\n",
        "        x = self.fc(x)\n",
        "\n",
        "        return F.softmax(x, dim=1)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5hTSKSxO6yxx"
      },
      "source": [
        "# From scratch\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ipdZR4DDBFpR"
      },
      "outputs": [],
      "source": [
        "class BasicBlock(nn.Module):\n",
        "    def __init__(self, in_channels, out_channels, stride, drop_rate):\n",
        "        super(BasicBlock, self).__init__()\n",
        "\n",
        "        self.bn1 = nn.BatchNorm2d(in_channels)\n",
        "        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=stride, padding=1, bias=False)\n",
        "\n",
        "        self.bn2 = nn.BatchNorm2d(out_channels)\n",
        "        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1, bias=False)\n",
        "\n",
        "        self.dropout = nn.Dropout(p=drop_rate) if drop_rate > 0 else nn.Identity()\n",
        "\n",
        "        if stride != 1 or in_channels != out_channels:\n",
        "            self.shortcut = nn.Sequential(\n",
        "                nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=stride, bias=False)\n",
        "            )\n",
        "        else:\n",
        "            self.shortcut = nn.Identity()\n",
        "\n",
        "    def forward(self, x):\n",
        "        out = F.relu(self.bn1(x))\n",
        "        out = self.conv1(out)\n",
        "\n",
        "        out = F.relu(self.bn2(out))\n",
        "        out = self.conv2(out)\n",
        "\n",
        "        out = self.dropout(out)\n",
        "\n",
        "        out += self.shortcut(x)\n",
        "        return out\n",
        "\n",
        "class WideResNet_or(nn.Module):\n",
        "    def __init__(self, depth, widen_factor, drop_rate, num_classes):\n",
        "        super(WideResNet, self).__init__()\n",
        "\n",
        "        n = (depth - 4) // 6\n",
        "        k = widen_factor\n",
        "\n",
        "        n_stages = [16, 16*k, 32*k, 64*k]\n",
        "\n",
        "        self.conv1 = nn.Conv2d(3, n_stages[0], kernel_size=3, stride=1, padding=1, bias=False)\n",
        "\n",
        "        self.layer1 = self._wide_layer(BasicBlock, n_stages[0], n_stages[1], n, 1, drop_rate)\n",
        "        self.layer2 = self._wide_layer(BasicBlock, n_stages[1], n_stages[2], n, 2, drop_rate)\n",
        "        self.layer3 = self._wide_layer(BasicBlock, n_stages[2], n_stages[3], n, 2, drop_rate)\n",
        "\n",
        "        self.bn = nn.BatchNorm2d(n_stages[3])\n",
        "        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))\n",
        "        self.fc = nn.Linear(n_stages[3], num_classes)\n",
        "\n",
        "    def _wide_layer(self, block, in_channels, out_channels, num_blocks, stride, drop_rate):\n",
        "        layers = []\n",
        "\n",
        "        for i in range(num_blocks):\n",
        "            layers.append(block(in_channels if i == 0 else out_channels,\n",
        "                                out_channels,\n",
        "                                stride if i == 0 else 1,\n",
        "                                drop_rate))\n",
        "\n",
        "        return nn.Sequential(*layers)\n",
        "\n",
        "    def forward(self, x):\n",
        "        x = self.conv1(x)\n",
        "\n",
        "        x = self.layer1(x)\n",
        "        x = self.layer2(x)\n",
        "        x = self.layer3(x)\n",
        "\n",
        "        x = F.relu(self.bn(x))\n",
        "\n",
        "        x = self.avgpool(x)\n",
        "        x = x.view(x.size(0), -1)\n",
        "\n",
        "        x = self.fc(x)\n",
        "\n",
        "        return F.log_softmax(x, dim=1)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Zm5afkjQ63jZ"
      },
      "source": [
        "# Sampler"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oC-3FmroAwbQ"
      },
      "source": [
        "# Sampler 2"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "HyLFdtJ1Au4l"
      },
      "outputs": [],
      "source": [
        "\n",
        "def accuracy(yhat, labels):\n",
        "    _, indices = yhat.max(1)\n",
        "    return (indices == labels).sum().data.item() / float(len(labels))\n",
        "\n",
        "class AverageCalculator():\n",
        "    def __init__(self):\n",
        "        self.reset()\n",
        "\n",
        "    def reset(self):\n",
        "        self.count = 0\n",
        "        self.sum = 0\n",
        "        self.avg = 0\n",
        "\n",
        "    def update(self, val, n=1):\n",
        "        assert(n > 0)\n",
        "        self.sum += val * n\n",
        "        self.count += n\n",
        "        self.avg = self.sum / float(self.count)\n",
        "\n",
        "\n",
        "\n",
        "class SVRG_k(Optimizer):\n",
        "    r\"\"\"Optimization class for calculating the gradient of one iteration.\n",
        "    Args:\n",
        "        params (iterable): iterable of parameters to optimize or dicts defining\n",
        "            parameter groups\n",
        "        lr (float): learning rate\n",
        "    \"\"\"\n",
        "    def __init__(self, params, lr, weight_decay=0):\n",
        "        print(\"Using optimizer: SVRG\")\n",
        "        self.u = None\n",
        "        if lr < 0.0:\n",
        "            raise ValueError(\"Invalid learning rate: {}\".format(lr))\n",
        "        if weight_decay < 0.0:\n",
        "            raise ValueError(\"Invalid weight decay: {}\".format(weight_decay))\n",
        "        defaults = dict(lr=lr, weight_decay=weight_decay)\n",
        "        super(SVRG_k, self).__init__(params, defaults)\n",
        "\n",
        "    def get_param_groups(self):\n",
        "            return self.param_groups\n",
        "\n",
        "    def set_u(self, new_u):\n",
        "        \"\"\"Set the mean gradient for the current epoch.\n",
        "        \"\"\"\n",
        "        if self.u is None:\n",
        "            self.u = copy.deepcopy(new_u)\n",
        "        for u_group, new_group in zip(self.u, new_u):\n",
        "            for u, new_u in zip(u_group['params'], new_group['params']):\n",
        "                u.grad = new_u.grad.clone()\n",
        "\n",
        "    def step(self, params):\n",
        "        \"\"\"Performs a single optimization step.\n",
        "        \"\"\"\n",
        "        for group, new_group, u_group in zip(self.param_groups, params, self.u):\n",
        "            weight_decay = group['weight_decay']\n",
        "\n",
        "            for p, q, u in zip(group['params'], new_group['params'], u_group['params']):\n",
        "                if p.grad is None:\n",
        "                    continue\n",
        "                if q.grad is None:\n",
        "                    continue\n",
        "                # core SVRG gradient update\n",
        "                new_d = p.grad.data - q.grad.data + u.grad.data\n",
        "                if weight_decay != 0:\n",
        "                    new_d.add_(weight_decay, p.data)\n",
        "                p.data.add_(-group['lr'], new_d)\n",
        "\n",
        "\n",
        "class SVRG_Snapshot(Optimizer):\n",
        "    r\"\"\"Optimization class for calculating the mean gradient (snapshot) of all samples.\n",
        "    Args:\n",
        "        params (iterable): iterable of parameters to optimize or dicts defining\n",
        "            parameter groups\n",
        "        lr (float): learning rate\n",
        "    \"\"\"\n",
        "    def __init__(self, params):\n",
        "        defaults = dict()\n",
        "        super(SVRG_Snapshot, self).__init__(params, defaults)\n",
        "\n",
        "    def get_param_groups(self):\n",
        "            return self.param_groups\n",
        "\n",
        "    def set_param_groups(self, new_params):\n",
        "        \"\"\"Copies the parameters from another optimizer.\n",
        "        \"\"\"\n",
        "        for group, new_group in zip(self.param_groups, new_params):\n",
        "            for p, q in zip(group['params'], new_group['params']):\n",
        "                  p.data[:] = q.data[:]\n",
        "\n",
        "\n",
        "\n",
        "\n",
        "def train_epoch_SVRG(model_k,\n",
        "                     model_snapshot,\n",
        "                     optimizer_k,\n",
        "                     optimizer_snapshot,\n",
        "                     train_loader,\n",
        "                     loss_fn,\n",
        "                     metric = accuracy,\n",
        "                     device = \"cuda\",\n",
        "                     data_transform = lambda x:x):\n",
        "    model_k.train()\n",
        "    model_snapshot.train()\n",
        "    loss = AverageCalculator()\n",
        "    av_metric = AverageCalculator()\n",
        "\n",
        "    # calculate the mean gradient\n",
        "    optimizer_snapshot.zero_grad()  # zero_grad outside for loop, accumulate gradient inside\n",
        "    for points, labels in train_loader:\n",
        "        points = points.to(device)\n",
        "        points = data_transform(points)\n",
        "\n",
        "        pred = model_snapshot(points)\n",
        "        labels = labels.to(device)\n",
        "        snapshot_loss = loss_fn(pred, labels) / len(train_loader)\n",
        "        snapshot_loss.backward()\n",
        "\n",
        "    # pass the current paramesters of optimizer_0 to optimizer_k\n",
        "    u = optimizer_snapshot.get_param_groups()\n",
        "    optimizer_k.set_u(u)\n",
        "\n",
        "    for points, labels in train_loader:\n",
        "        points = points.to(device)\n",
        "        points = data_transform(points)\n",
        "\n",
        "        pred_1 = model_k(points)\n",
        "        labels = labels.to(device)\n",
        "        loss_iter = loss_fn(pred_1, labels)\n",
        "\n",
        "        # optimization\n",
        "        optimizer_k.zero_grad()\n",
        "        loss_iter.backward()\n",
        "\n",
        "        pred_2 = model_snapshot(points)\n",
        "        loss_2 = loss_fn(pred_2, labels)\n",
        "\n",
        "        optimizer_snapshot.zero_grad()\n",
        "        loss_2.backward()\n",
        "\n",
        "        optimizer_k.step(optimizer_snapshot.get_param_groups())\n",
        "\n",
        "        loss.update(loss_iter.data.item())\n",
        "        av_metric.update(metric(pred_1, labels))\n",
        "\n",
        "    # update the snapshot\n",
        "    optimizer_snapshot.set_param_groups(optimizer_k.get_param_groups())\n",
        "\n",
        "    return loss.avg, av_metric.avg\n",
        "\n",
        "\n",
        "\n",
        "def test_epoch_SVRG(model, test_loader, loss_fn, metric = accuracy, device = \"cuda\", data_transform = lambda x:x):\n",
        "    \"\"\"One epoch of validation\n",
        "    \"\"\"\n",
        "    model.eval()\n",
        "    loss = AverageCalculator()\n",
        "    av_metric = AverageCalculator()\n",
        "    with torch.no_grad():\n",
        "        for points, labels in test_loader:\n",
        "            points = points.to(device)\n",
        "            points = data_transform(points)\n",
        "            pred = model(points)\n",
        "            labels = labels.to(device)\n",
        "\n",
        "            loss_iter = loss_fn(pred, labels)\n",
        "            av_metric.update(metric(pred, labels))\n",
        "            loss.update(loss_iter.data.item())\n",
        "\n",
        "\n",
        "    return loss.avg, av_metric.avg"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "24WKmKMo669Q"
      },
      "source": [
        "# Training"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "X9edbl0MCDPf"
      },
      "outputs": [],
      "source": [
        "import torchvision.transforms as transforms\n",
        "from torchvision.datasets import CIFAR10\n",
        "from torch.utils.data import DataLoader"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bdyp2Tx7-R5m"
      },
      "outputs": [],
      "source": [
        "input_shape = (3, 32, 32)  # Assuming CIFAR-10 dataset\n",
        "output_size = 10\n",
        "\n",
        "# Create an instance of the WideResNet model with depth=28 and widen_factor=10\n",
        "model = WideResNet(depth=16, widen_factor=4, drop_rate=0.0, input_shape=input_shape, output_size=output_size)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "1sQs4qZGEHKL",
        "outputId": "a60862b1-112d-44b4-b0ef-9fdc044d6e48"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Number of weights: 3101242\n"
          ]
        }
      ],
      "source": [
        "print('Number of weights:', np.sum([np.prod(p.shape) for p in model.parameters()]))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "nak0PhUGB_Ih",
        "outputId": "dc6c79c5-6120-4e17-873b-178b177156c5"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Downloading https://www.cs.toronto.edu/~kriz/cifar-10-python.tar.gz to ./data/cifar-10-python.tar.gz\n"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "100%|██████████| 170498071/170498071 [00:03<00:00, 49573637.57it/s]\n"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Extracting ./data/cifar-10-python.tar.gz to ./data\n",
            "Files already downloaded and verified\n"
          ]
        }
      ],
      "source": [
        "mean, std = [0.5, 0.5, 0.5] , [0.5, 0.5, 0.5]\n",
        "train_transform = transforms.Compose([\n",
        "    transforms.RandomHorizontalFlip(),\n",
        "    transforms.RandomRotation(20),\n",
        "    transforms.ToTensor(),\n",
        "    transforms.Normalize(mean, std)\n",
        "])\n",
        "test_transform = transforms.Compose([\n",
        "    transforms.ToTensor(),\n",
        "    transforms.Normalize(mean, std)\n",
        "])\n",
        "\n",
        "train_set = CIFAR10(root='./data', train=True, download=True, transform=train_transform)\n",
        "train_loader = DataLoader(train_set, batch_size=32, shuffle=True)\n",
        "\n",
        "test_set = CIFAR10(root='./data', train=False, download=True, transform=test_transform)\n",
        "test_loader = DataLoader(test_set, batch_size=32, shuffle=False)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "fMby6KTfCLkW",
        "outputId": "65cc4f01-fbcd-41d6-b690-88fbcdb5201f"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Train size 50000\n",
            "Test size 10000\n"
          ]
        }
      ],
      "source": [
        "classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck')\n",
        "\n",
        "print('Train size', len(train_set))\n",
        "print('Test size', len(test_set))\n",
        "\n",
        "n_epochs = 10"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "y-M-_7KqEUDO",
        "outputId": "06db45d6-86dc-4010-aeaf-1ec1867609c4"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Using optimizer: SVRG\n"
          ]
        }
      ],
      "source": [
        "device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')\n",
        "model = model.to(device)\n",
        "model_snapshot = CNN().to(device)\n",
        "\n",
        "loss_fn = nn.CrossEntropyLoss()\n",
        "\n",
        "optimizer = SVRG_k(model.parameters(), lr=0.02, weight_decay = 0.02)\n",
        "optimizer_snapshot = SVRG_Snapshot(model_snapshot.parameters())"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "xHwDIzS5B4XV",
        "outputId": "3cdb87d2-d1a2-4113-8217-97a4ac0ffce6"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "  0%|          | 0/10 [25:26<?, ?it/s]\n"
          ]
        },
        {
          "ename": "RuntimeError",
          "evalue": "running_mean should contain 16 elements not 64",
          "output_type": "error",
          "traceback": [
            "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
            "\u001b[0;31mRuntimeError\u001b[0m                              Traceback (most recent call last)",
            "\u001b[0;32m<ipython-input-47-263f2d055c98>\u001b[0m in \u001b[0;36m<cell line: 1>\u001b[0;34m()\u001b[0m\n\u001b[1;32m      1\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mepoch\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mtqdm\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mrange\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mn_epochs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 2\u001b[0;31m     \u001b[0mtrain_loss\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtrain_acc\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtrain_epoch_SVRG\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmodel\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmodel_snapshot\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0moptimizer\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0moptimizer_snapshot\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtrain_loader\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mloss_fn\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0mdevice\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m'cpu'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m      3\u001b[0m     \u001b[0mtest_loss\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtest_acc\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mtest_epoch_SVRG\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mmodel\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtest_loader\u001b[0m\u001b[0;34m,\u001b[0m  \u001b[0mloss_fn\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m      4\u001b[0m     print(f'[Epoch {epoch + 1}] train loss: {train_loss:.3f}; train acc: {train_acc:.2f}; ' + \n\u001b[1;32m      5\u001b[0m           f'test loss: {test_loss:.3f}; test acc: {test_acc:.2f}')\n",
            "\u001b[0;32m<ipython-input-39-d98beb2f908b>\u001b[0m in \u001b[0;36mtrain_epoch_SVRG\u001b[0;34m(model_k, model_snapshot, optimizer_k, optimizer_snapshot, train_loader, loss_fn, metric, device, data_transform)\u001b[0m\n\u001b[1;32m    124\u001b[0m         \u001b[0mpoints\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mdata_transform\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mpoints\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    125\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 126\u001b[0;31m         \u001b[0mpred_1\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mmodel_k\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mpoints\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    127\u001b[0m         \u001b[0mlabels\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mlabels\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mto\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mdevice\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    128\u001b[0m         \u001b[0mloss_iter\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mloss_fn\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mpred_1\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mlabels\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_wrapped_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1516\u001b[0m             \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_compiled_call_impl\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m  \u001b[0;31m# type: ignore[misc]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1517\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1518\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_call_impl\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1519\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1520\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0m_call_impl\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1525\u001b[0m                 \u001b[0;32mor\u001b[0m \u001b[0m_global_backward_pre_hooks\u001b[0m \u001b[0;32mor\u001b[0m \u001b[0m_global_backward_hooks\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1526\u001b[0m                 or _global_forward_hooks or _global_forward_pre_hooks):\n\u001b[0;32m-> 1527\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mforward_call\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1528\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1529\u001b[0m         \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m<ipython-input-36-8f784fd73711>\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, x)\u001b[0m\n\u001b[1;32m     40\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mforward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     41\u001b[0m         \u001b[0mx\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mgroup0\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 42\u001b[0;31m         \u001b[0mx\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mgroup1\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     43\u001b[0m         \u001b[0mx\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mgroup2\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     44\u001b[0m         \u001b[0mx\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mgroup3\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mx\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_wrapped_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1516\u001b[0m             \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_compiled_call_impl\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m  \u001b[0;31m# type: ignore[misc]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1517\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1518\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_call_impl\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1519\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1520\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0m_call_impl\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1525\u001b[0m                 \u001b[0;32mor\u001b[0m \u001b[0m_global_backward_pre_hooks\u001b[0m \u001b[0;32mor\u001b[0m \u001b[0m_global_backward_hooks\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1526\u001b[0m                 or _global_forward_hooks or _global_forward_pre_hooks):\n\u001b[0;32m-> 1527\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mforward_call\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1528\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1529\u001b[0m         \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/torch/nn/modules/container.py\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, input)\u001b[0m\n\u001b[1;32m    213\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mforward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minput\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    214\u001b[0m         \u001b[0;32mfor\u001b[0m \u001b[0mmodule\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 215\u001b[0;31m             \u001b[0minput\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mmodule\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    216\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0minput\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    217\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_wrapped_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1516\u001b[0m             \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_compiled_call_impl\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m  \u001b[0;31m# type: ignore[misc]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1517\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1518\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_call_impl\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1519\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1520\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0m_call_impl\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1525\u001b[0m                 \u001b[0;32mor\u001b[0m \u001b[0m_global_backward_pre_hooks\u001b[0m \u001b[0;32mor\u001b[0m \u001b[0m_global_backward_hooks\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1526\u001b[0m                 or _global_forward_hooks or _global_forward_pre_hooks):\n\u001b[0;32m-> 1527\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mforward_call\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1528\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1529\u001b[0m         \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/torch/nn/modules/container.py\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, input)\u001b[0m\n\u001b[1;32m    213\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mforward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minput\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    214\u001b[0m         \u001b[0;32mfor\u001b[0m \u001b[0mmodule\u001b[0m \u001b[0;32min\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 215\u001b[0;31m             \u001b[0minput\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mmodule\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    216\u001b[0m         \u001b[0;32mreturn\u001b[0m \u001b[0minput\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    217\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_wrapped_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1516\u001b[0m             \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_compiled_call_impl\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m  \u001b[0;31m# type: ignore[misc]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1517\u001b[0m         \u001b[0;32melse\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 1518\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_call_impl\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1519\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1520\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0m_call_impl\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py\u001b[0m in \u001b[0;36m_call_impl\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1525\u001b[0m                 \u001b[0;32mor\u001b[0m \u001b[0m_global_backward_pre_hooks\u001b[0m \u001b[0;32mor\u001b[0m \u001b[0m_global_backward_hooks\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1526\u001b[0m                 or _global_forward_hooks or _global_forward_pre_hooks):\n\u001b[0;32m-> 1527\u001b[0;31m             \u001b[0;32mreturn\u001b[0m \u001b[0mforward_call\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m*\u001b[0m\u001b[0margs\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m**\u001b[0m\u001b[0mkwargs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m   1528\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   1529\u001b[0m         \u001b[0;32mtry\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/torch/nn/modules/batchnorm.py\u001b[0m in \u001b[0;36mforward\u001b[0;34m(self, input)\u001b[0m\n\u001b[1;32m    169\u001b[0m         \u001b[0mused\u001b[0m \u001b[0;32mfor\u001b[0m \u001b[0mnormalization\u001b[0m \u001b[0;34m(\u001b[0m\u001b[0mi\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0me\u001b[0m\u001b[0;34m.\u001b[0m \u001b[0;32min\u001b[0m \u001b[0meval\u001b[0m \u001b[0mmode\u001b[0m \u001b[0mwhen\u001b[0m \u001b[0mbuffers\u001b[0m \u001b[0mare\u001b[0m \u001b[0;32mnot\u001b[0m \u001b[0;32mNone\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    170\u001b[0m         \"\"\"\n\u001b[0;32m--> 171\u001b[0;31m         return F.batch_norm(\n\u001b[0m\u001b[1;32m    172\u001b[0m             \u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    173\u001b[0m             \u001b[0;31m# If buffers are not to be tracked, ensure that they won't be updated\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
            "\u001b[0;32m/usr/local/lib/python3.10/dist-packages/torch/nn/functional.py\u001b[0m in \u001b[0;36mbatch_norm\u001b[0;34m(input, running_mean, running_var, weight, bias, training, momentum, eps)\u001b[0m\n\u001b[1;32m   2476\u001b[0m         \u001b[0m_verify_batch_size\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0minput\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msize\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   2477\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m-> 2478\u001b[0;31m     return torch.batch_norm(\n\u001b[0m\u001b[1;32m   2479\u001b[0m         \u001b[0minput\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mweight\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mbias\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mrunning_mean\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mrunning_var\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtraining\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mmomentum\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0meps\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbackends\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mcudnn\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0menabled\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m   2480\u001b[0m     )\n",
            "\u001b[0;31mRuntimeError\u001b[0m: running_mean should contain 16 elements not 64"
          ]
        }
      ],
      "source": [
        "for epoch in tqdm(range(n_epochs)):\n",
        "    train_loss, train_acc = train_epoch_SVRG(model, model_snapshot, optimizer, optimizer_snapshot, train_loader, loss_fn,device='cpu')\n",
        "    test_loss, test_acc = test_epoch_SVRG(model, test_loader,  loss_fn)\n",
        "    print(f'[Epoch {epoch + 1}] train loss: {train_loss:.3f}; train acc: {train_acc:.2f}; ' +\n",
        "          f'test loss: {test_loss:.3f}; test acc: {test_acc:.2f}')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "M66daLAWCP7Z"
      },
      "outputs": [],
      "source": []
    }
  ],
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}