MNIST.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "fiSHwkvZumRd"
   },
   "source": [
    "# Experiments on MNIST"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from optimizers import *\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import time\n",
    "%matplotlib inline\n",
    "\n",
    "import torch\n",
    "from torch import nn\n",
    "from torch import optim\n",
    "from torch.autograd import Variable\n",
    "from torch.utils.data.sampler import SubsetRandomSampler\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "\n",
    "import torchvision\n",
    "import torchvision.transforms as transforms"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 34
    },
    "colab_type": "code",
    "id": "p2Mo1UdyumT9",
    "outputId": "f14053f0-1237-42a4-f9ee-f5c16460b496"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "device(type='cpu')"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
    "device"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "file_path = \"./MNIST\"\n",
    "batch_size = 4"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "W6Lk5gvGu1Yc"
   },
   "source": [
    "## Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "_rk3Q_Jku2by"
   },
   "outputs": [],
   "source": [
    "transform = transforms.Compose([transforms.ToTensor(), \n",
    "                                transforms.Normalize((0.5,), (0.5,))])\n",
    "trainset = torchvision.datasets.MNIST(root='../data', \n",
    "                                      train=True, download=True, transform=transform)\n",
    "valid_dataset = torchvision.datasets.MNIST(root='../data', \n",
    "                                           train=True, download=True, transform=transform)\n",
    "testset = torchvision.datasets.MNIST(root='../data', \n",
    "                                     train=False, download=True, transform=transform)\n",
    "testloader = torch.utils.data.DataLoader(testset, batch_size=batch_size, shuffle=False, num_workers=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "bYLS8LcHWv7n"
   },
   "outputs": [],
   "source": [
    "valid_size=0.15\n",
    "num_train = len(trainset)\n",
    "indices = list(range(num_train))\n",
    "split = int(np.floor(valid_size * num_train))\n",
    "\n",
    "np.random.seed(42)\n",
    "np.random.shuffle(indices)\n",
    "\n",
    "train_idx, valid_idx = indices[split:], indices[:split]\n",
    "\n",
    "train_sampler = SubsetRandomSampler(train_idx)\n",
    "valid_sampler = SubsetRandomSampler(valid_idx)\n",
    "\n",
    "trainloader = torch.utils.data.DataLoader(trainset, \n",
    "               batch_size=batch_size, sampler=train_sampler,\n",
    "               num_workers=2)\n",
    "\n",
    "validloader = torch.utils.data.DataLoader(valid_dataset, \n",
    "               batch_size=batch_size, sampler=valid_sampler,\n",
    "               num_workers=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 190
    },
    "colab_type": "code",
    "id": "7j9cjtv5u57l",
    "outputId": "2d6a7e4b-aa09-49fd-9dbd-d0e3f65ff814"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Clipping input data to the valid range for imshow with RGB data ([0..1] for floats or [0..255] for integers).\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Labels:  tensor([7, 1, 8, 2])\n",
      "Batch shape:  torch.Size([4, 1, 28, 28])\n"
     ]
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAXYAAAB6CAYAAACr63iqAAAABHNCSVQICAgIfAhkiAAAAAlwSFlzAAALEgAACxIB0t1+/AAAADl0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uIDIuMi4zLCBodHRwOi8vbWF0cGxvdGxpYi5vcmcvIxREBQAAD4JJREFUeJzt3X+sFfWZx/H34+VXbRPlFkQElWsWBNRFiTFgVzQqLtamNGoQg12MKFHrDwzRXkGDvxJZlYpr1iopbOkGZd0Li4R1dwX8lQbKilti4d6KWBRBBAyI3TWhoM/+MXOGQe7hnDPnx5wz9/NKyH3OzJwzz9y5fv2eZ77zHXN3REQkO45LOwEREaksNewiIhmjhl1EJGPUsIuIZIwadhGRjFHDLiKSMWrYRUQypqyG3czGmdn7ZrbFzForlZSIiCRnSW9QMrMmYDMwFtgOvANc7+7tlUtPRERK1a2M914AbHH3PwGY2WJgPJC3YTcz3eYqIlK6z929b7Ebl1OKGQB8Enu9PVx2BDObambrzWx9GfsSEenKPi5l43J67EVx93nAPFCPXUSkFsrpse8ATo29HhguExGRFJXTsL8DDDazFjPrAUwEllcmLRERSSpxKcbdD5nZHcB/AU3AAnffVLHMREQkkcTDHRPtTDV2EZEk3nX384vdWHeeiohkjBp2EZGMUcMuIpIxathFRDJGDbuISMaoYRcRyRg17CIiGaOGXUQkY9Swi4hkTNVndxRJ6uGHH47iWbNmpZhJ7W3dujWKW1paUsykcZxzzjlRvGHDhmNu29TUVO10UqUeu4hIxqhhFxHJGJVipK4sWbIkij/66KP0EknBjh2HH2fQp0+fKD5w4MBR23711VdRPGTIkCjes2dPlbKrL539TuIOHTpU9Pu7dTvcDC5btgyAa665pozs0qceu4hIxqhhFxHJGM3HLnVl//79UXzCCSekmEntFSovxB133OE+WWtraxTPmTOnojnVk1J+P6WIl2I6K+HMmDEjilP8/Wo+dhGRrkwNu4hIxqgUI3Vh06bgcbk33XRTtGzdunVppZOKtWvXRvHIkSOPuW28FPPNN99E8cCBA6M4ayNkSinFPPHEE1G8atWqo9a3tbVFcXNzcxQXGk1z1113RfELL7xQdD4VUNlSjJktMLPdZrYxtqzZzFaa2Qfhz95JsxURkcoq2GM3szHA/wK/cfezw2VPAHvdfbaZtQK93f3nBXemHrvEXHTRRVG8dOlSAPr27ZtWOqmbNGlSFC9YsOCY2+brsW/evDmK47fYZ0Hv3of7jxMmTABg6NCh0bJHHnkkivft21f053799ddRXKjHHnf77bcDMH/+/KLfU4bK9tjd/W1g77cWjwcWhvFC4CdFpyciIlWV9M7Tfu6+M4w/A/rl29DMpgJTE+5HRERKVNTFUzMbBKyIlWK+cPcTY+v3uXvBOnvSUszYsWMBePXVV6Nl3bt3T/JRiXVWIjjppJOi+MUXX4ziESNG1CSnRjRs2LAo3rgxumzDhRdeCHS9C6ZxDz74YBQ/8MADUbxixYoozt3qfvDgwWhZvBQT17Nnz0qnmHn9+gV91G3bthX9nhr9nmsyjn2XmfUHCH/uTvg5IiJSYUkb9uXA5DCeDLxSmXRERKRcxYyKeQm4BOgD7AJmAcuAl4HTgI+BCe7+7QusnX1WolJM//79gfyz/eVGVMDh0QT5vp7GXXvttQDcfPPN0bJi3peTb2SCvgLnlxuvDrBmzZoovuWWW9JIp67kG6fd2d+TSjHVVcqY+aeeeiqKZ86cWY10oMRSTMGLp+5+fZ5VlxWdkoiI1IymFBARyZiGmlKgmK9HufJIJUoqSd+3ZcsWAM4666yiPyvLCo2Ega49Gubee+8F4LHHHut0vUoxtZd0Jskq/s41u6OISFfWUD32+JPF29vbo3jQoEFRXG6Pffbs2VF8+eWXR/HevcG14fiFkqlTD993dfXVVx/1ueoxBTo6OqK4V69eUdzS0pJGOnXnrbfeAmDUqFGdrp87d24UT5s2DSjuW2Zu8qurrrqqInl2VaX03tVjFxGRqlDDLiKSMQ1Viqk38WkGtm/fftR6lWIC8dnz4uW0ruyZZ56J4ltvvbXk95dywV9/h+VRKUZERFKnhl1EJGOSTtsrHPl1OB7fd999aaRTd5599lkg/1QQXVnuIQ1Jdet2+D/dfA+HuPHGG8vaR1cWLx/GJb3npdbUYxcRyRg17CIiGaNSTBniX8XyxV1ZrtwQnzpAAvEHxRS6QakzQ4YMieKJEydGcfwBHc899xwAixYtSpxnvUt6639cZ2WtUp59Wo/UYxcRyRj12MuwZ8+etFOoO/Hx2bnHBXblCb6KkXv0Xb4ee3xysEcfffSo9fFl8THxffr0AWD69OnRsjlz5pSXbJ2oRE+9GuJ5pXn/gHrsIiIZo4ZdRCRjVIqRsp155plRfMcdd0SxLpoml5tNFDovv+QTn300N1Pp8OHDK5dYiuq1/FKPCvbYzexUM3vDzNrNbJOZ3R0ubzazlWb2Qfizd/XTFRGRQoopxRwCprv7cGAU8DMzGw60AqvdfTCwOnwtIiIpK+Zh1juBnWH8ZzPrAAYA44FLws0WAm8CP69KlnUqPrtjV/b6669HcVtbWxRrNExyuccrlurpp5+O4lwp5oYbboiWTZkypbzEaixf+SX3dxYv/RXy5ptvRvHQoUPLyqsYo0ePjuK1a9dWfX9xJdXYzWwQcB6wDugXNvoAnwH98rxnKjC1s3UiIlJ5RY+KMbPvAUuAae7+ZXydB5O6dzrXurvPc/fzS5lLWEREkiuqx25m3Qka9UXuvjRcvMvM+rv7TjPrD+yuVpL1Kt/sjvGbdLIqPhLm5JNPjuLrrrsujXQypxKlgjVr1gBHjk6K36AUv3Gp0cybNw+Affv21XS/Dz30UBQ//vjjNd13KYoZFWPAfKDD3X8RW7UcmBzGk4FXKp+eiIiUqpge+w+AnwJ/MLMN4bIZwGzgZTObAnwMTKhOivUr3qvqahN/5XqDAMuXL08xk8b35JNPAkfO0X7KKadEcfyC6D333FP0555xxhlHLWtvb0+SYmrit+XHL6S2tgaD8Hr06FHwM3JTMlTiW9Dzzz9f9mfUQjGjYn4LWJ7Vl1U2HRERKZemFBARyRgLBrTUaGdmtdtZDRw8eDCKP/300yg+/fTT00inJnIXTeNf6ZuamtJKJ1MuvvjiKH7ttdeOuW2+mQNvu+22KJ47dy4AX355eBBbVu69qOT0Ap3Nx567OAtw5513VmxfZXi3lJGF6rGLiGSMGnYRkYzR7I4VsmrVqrRTqInm5magcKlASpd7RB4cWWrZunVrFOdGy8TLgIVGZM2YMaNSKdaN3O+nEiWZzZs3R/GwYcPK/rx6oB67iEjGqGEXEckYlWKkJLlZ6q644oqUM+k6WlpajloWL8Xkc9pppwGwa9euiudUL+Ilq3gZZcOGDZ1t3un7skg9dhGRjFGPvQzx8a+ff/55ipnUnsaup6t79+5pp1B3Ojo6ojjrPfJC1GMXEckYNewiIhmjKQXKEL+Apa/GIlJFmlJARKQrU8MuIpIxGhVTovhDD0RE6pF67CIiGaOGXUQkYwqWYsysF/A20DPcvs3dZ5lZC7AY+D7wLvBTd/9LNZOtB2PGjIni+KxwIiL1opge+wHgUncfAZwLjDOzUcDfA0+7+18B+4Ap1UtTRESKVdI4djM7HvgtcBvw78DJ7n7IzEYDD7n73xZ4f6bGsYuI1Ejlx7GbWZOZbQB2AyuBD4Ev3P1QuMl2YECpmYqISOUV1bC7+9fufi4wELgAGFrsDsxsqpmtN7P1CXMUEZESlDQqxt2/AN4ARgMnmlnu4utAYEee98xz9/NL+RohIiLJFWzYzayvmZ0Yxt8BxgIdBA38teFmk4FXqpWkiIgUr5g7T/sDC82sieB/BC+7+wozawcWm9ljwO+B+VXMU0REilTr2R33AP8HZPWpFH3QsTUiHVtj6krHdrq79y32zTVt2AHMbH1W6+06tsakY2tMOrb8NKWAiEjGqGEXEcmYNBr2eSnss1Z0bI1Jx9aYdGx51LzGLiIi1aVSjIhIxqhhFxHJmJo27GY2zszeN7MtZtZay31XmpmdamZvmFm7mW0ys7vD5c1mttLMPgh/9k471yTCid9+b2YrwtctZrYuPHf/YmY90s4xCTM70czazOyPZtZhZqMzdM7uCf8WN5rZS2bWq1HPm5ktMLPdZrYxtqzT82SBfwiP8T0zG5le5oXlObYnw7/J98zs33J3+4fr7g+P7X0zO+YMujk1a9jDO1f/EbgSGA5cb2bDa7X/KjgETHf34cAo4Gfh8bQCq919MLA6fN2I7iaYOiInK/PvPwP8p7sPBUYQHGPDnzMzGwDcBZzv7mcDTcBEGve8/RoY961l+c7TlcDg8N9U4Jc1yjGpX3P0sa0Eznb3vwY2A/cDhG3KROCs8D3PhW3pMdWyx34BsMXd/xQ+aWkxML6G+68od9/p7v8Txn8maCAGEBzTwnCzhcBP0skwOTMbCFwF/Cp8bcClQFu4SaMe1wnAGMLpL9z9L+HEdg1/zkLdgO+Ek/MdD+ykQc+bu78N7P3W4nznaTzwGw/8jmCCwv61ybR0nR2bu78Wmwb9dwQTK0JwbIvd/YC7bwW2ELSlx1TLhn0A8EnsdWbmcDezQcB5wDqgn7vvDFd9BvRLKa1yzAXuA74JX3+fbMy/3wLsAf4pLDP9ysy+SwbOmbvvAJ4CthE06PsJHlmZhfOWk+88Za1tuQn4jzBOdGy6eFomM/sesASY5u5fxtd5MJa0ocaTmtmPgN3u/m7auVRBN2Ak8Et3P49g3qIjyi6NeM4AwnrzeIL/eZ0CfJejv+5nRqOep0LMbCZBmXdROZ9Ty4Z9B3Bq7HXeOdwbhZl1J2jUF7n70nDxrtzXwPDn7rTyS+gHwI/N7COCctmlBHXpoubfr3Pbge3uvi583UbQ0Df6OQO4HNjq7nvc/SCwlOBcZuG85eQ7T5loW8zsRuBHwCQ/fINRomOrZcP+DjA4vErfg+CCwPIa7r+iwrrzfKDD3X8RW7WcYH56aMB56t39fncf6O6DCM7R6+4+iQzMv+/unwGfmNmZ4aLLgHYa/JyFtgGjzOz48G8zd2wNf95i8p2n5cDfhaNjRgH7YyWbhmBm4wjKnz92969iq5YDE82sp5m1EFwg/u+CH+juNfsH/JDgiu+HwMxa7rsKx/I3BF8F3wM2hP9+SFCPXg18AKwCmtPOtYxjvARYEcZnhH9QW4B/BXqmnV/CYzoXWB+et2VA76ycM+Bh4I/ARuCfgZ6Net6AlwiuFRwk+KY1Jd95AoxgxN2HwB8IRgalfgwlHtsWglp6ri15Prb9zPDY3geuLGYfmlJARCRjdPFURCRj1LCLiGSMGnYRkYxRwy4ikjFq2EVEMkYNu4hIxqhhFxHJmP8HVxSpg5BPY5YAAAAASUVORK5CYII=\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "def show_batch(batch):\n",
    "    im = torchvision.utils.make_grid(batch)\n",
    "    plt.imshow(np.transpose(im.numpy(), (1, 2, 0)))\n",
    "    \n",
    "dataiter = iter(trainloader)\n",
    "images, labels = dataiter.next()\n",
    "\n",
    "print('Labels: ', labels)\n",
    "print('Batch shape: ', images.size())\n",
    "show_batch(images)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "THk9uDZku-gM"
   },
   "source": [
    "## Models"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "pnS_0OICbHF4"
   },
   "source": [
    "### Logistic Regression"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "H6micuYGbGOL"
   },
   "outputs": [],
   "source": [
    "class LR(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(LR, self).__init__()\n",
    "        self.linear1 = nn.Linear(28*28, 10)\n",
    "\n",
    "    def forward(self, x):\n",
    "        y_pred = F.log_softmax(self.linear1(x.view(batch_size, -1)), -1)\n",
    "        return y_pred"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "A4rjRazJa-S6"
   },
   "source": [
    "###   Two-layer neural network (fully connected)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "uh3HnmUmu96Q"
   },
   "outputs": [],
   "source": [
    "class FC(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(FC, self).__init__()\n",
    "        self.linear1 = nn.Linear(28*28, 256)\n",
    "        self.linear2 = nn.Linear(256, 10)\n",
    "\n",
    "    def forward(self, x):\n",
    "        h_relu = F.relu(self.linear1(x.view(batch_size, -1)))\n",
    "        y_pred = self.linear2(h_relu)\n",
    "        return y_pred"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "-SSTdIH_vKV5"
   },
   "source": [
    "## Train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "nMttjLlZrbcd"
   },
   "outputs": [],
   "source": [
    "def time_since(since):\n",
    "    s = time.time() - since\n",
    "    m = math.floor(s / 60)\n",
    "    s -= m * 60\n",
    "    return '%dm %ds' % (m, s)\n",
    "\n",
    "\n",
    "def model_step(model, optimizer, criterion, inputs, labels):\n",
    "    outputs = model(inputs)\n",
    "    loss = criterion(outputs, labels)\n",
    "\n",
    "    if model.training:\n",
    "        optimizer.zero_grad()\n",
    "        loss.backward(retain_graph=True)\n",
    "    if optimizer.__class__.__name__ != 'SUG':\n",
    "        optimizer.step()\n",
    "    else:\n",
    "        def closure():\n",
    "            optimizer.zero_grad()\n",
    "            upd_outputs = model(inputs)\n",
    "            upd_loss = criterion(upd_outputs, labels)\n",
    "            upd_loss.backward()\n",
    "            return upd_loss\n",
    "\n",
    "        optimizer.step(loss, closure)\n",
    "\n",
    "    return loss.item()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "MMgATVQnvIiy"
   },
   "outputs": [],
   "source": [
    "def train(model, trainloader, criterion, optimizer, n_epochs=2, validloader=None, eps=1e-5, print_every=1):\n",
    "    tr_loss, val_loss, lips, times, grad, acc = ([] for i in range(6))\n",
    "    start_time = time.time()\n",
    "    model.to(device=device)\n",
    "    for ep in range(n_epochs):\n",
    "        model.train()\n",
    "        i = 0\n",
    "        for i, data in enumerate(trainloader):\n",
    "            inputs, labels = data\n",
    "            inputs, labels = Variable(inputs).to(device=device), Variable(labels).to(device=device)\n",
    "\n",
    "            tr_loss.append(model_step(model, optimizer, criterion, inputs, labels))\n",
    "            if optimizer.__class__.__name__ == 'SUG':\n",
    "                lips.append(optimizer.get_lipsitz_const())\n",
    "                grad.append(optimizer.get_sq_grad)\n",
    "        times.append(time_since(start_time))\n",
    "        if ep % print_every == 0:\n",
    "            print(\"Epoch {}, training loss {}, time passed {}\".format(ep, sum(tr_loss[-i:]) / i, time_since(start_time)))\n",
    "\n",
    "        if validloader is None:\n",
    "            continue\n",
    "        model.zero_grad()\n",
    "        model.eval()\n",
    "        j = 0\n",
    "        for j, data in enumerate(validloader):\n",
    "            inputs, labels = data\n",
    "            inputs, labels = inputs.to(device=device), labels.to(device=device)\n",
    "            val_loss.append(model_step(model, optimizer, criterion, inputs, labels))\n",
    "        if ep % print_every == 0:\n",
    "            print(\"Validation loss {}\".format(sum(val_loss[-j:]) / j))\n",
    "        \n",
    "    return tr_loss, times, val_loss, lips, grad"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "Sz4svVB3vNyZ"
   },
   "outputs": [],
   "source": [
    "print_every = 2\n",
    "n_epochs = 6\n",
    "tr_loss = {}\n",
    "tr_loss['accSGD'] = {}\n",
    "tr_loss['Adam'] = {}\n",
    "tr_loss['amsgrad'] = {}\n",
    "tr_loss['sug'] = {}\n",
    "tr_loss['A2GradInc'] = {}\n",
    "tr_loss['A2GradUni'] = {}\n",
    "tr_loss['A2GradExp'] = {}\n",
    "val_loss = {}\n",
    "val_loss['accSGD'] = {}\n",
    "val_loss['Adam'] = {}\n",
    "val_loss['amsgrad'] = {}\n",
    "val_loss['sug'] = {}\n",
    "val_loss['A2GradInc'] = {}\n",
    "val_loss['A2GradUni'] = {}\n",
    "val_loss['A2GradExp'] = {}\n",
    "lrs = [0.001]#[0.1, 0.01, 0.001]\n",
    "criterion = nn.CrossEntropyLoss(reduction=\"mean\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "B6jCFM66X7ql"
   },
   "outputs": [],
   "source": [
    "def concat_states(state1, state2):\n",
    "    states = {\n",
    "            'epoch': state1['epoch'] + state2['epoch'],\n",
    "            'state_dict': state2['state_dict'],\n",
    "            'optimizer': state2['optimizer'],\n",
    "            'tr_loss' : state1['tr_loss'] + state2['tr_loss'],\n",
    "            'val_loss' : state1['val_loss'] + state2['val_loss'],\n",
    "            'lips' : state1['lips'] + state2['lips'],\n",
    "            'grad' : state1['grad'] + state2['grad'],\n",
    "            #'times' : state1['times'] + list(map(lambda x: x + state1['times'][-1],state2['times']))\n",
    "             'times' : state1['times'] + state2['times']\n",
    "             }\n",
    "    return states"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "98EHtawSzoTM"
   },
   "source": [
    "### LR"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Adam "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 374
    },
    "colab_type": "code",
    "id": "hTMWs2dYznyO",
    "outputId": "01e33386-4143-4fe1-ead2-ef3ccd13cea8"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Adam  lr=0.001:\n",
      "Epoch 0, training loss 0.4471673777217208, time passed 0m 17s\n",
      "Validation loss 0.48481341447177106\n",
      "Epoch 2, training loss 0.39360389989321554, time passed 0m 54s\n",
      "Validation loss 0.4692441363024516\n",
      "Epoch 4, training loss 0.38341145817109673, time passed 1m 35s\n",
      "Validation loss 0.5846125006058818\n"
     ]
    }
   ],
   "source": [
    "for lr in lrs:\n",
    "    model = LR()\n",
    "    print(\"Adam  lr={}:\".format(lr))\n",
    "    optimizer = optim.Adam(model.parameters(), lr=lr)\n",
    "    tr_loss['Adam'][lr], times, val_loss['Adam'][lr], lips, grad = train(model, trainloader, criterion, \n",
    "                                                                         optimizer, n_epochs=n_epochs, \n",
    "                                                                         print_every=print_every, \n",
    "                                                                         validloader=validloader)\n",
    "    states = {\n",
    "            'epoch': n_epochs,\n",
    "            'state_dict': model.state_dict(),\n",
    "            'optimizer': optimizer.state_dict(),\n",
    "            'tr_loss' : tr_loss['Adam'][lr],\n",
    "            'val_loss' : val_loss['Adam'][lr],\n",
    "            'lips' : lips,\n",
    "            'grad' : grad,\n",
    "            'times' : times\n",
    "             }\n",
    "    torch.save(states, './MNIST/LR_Adam_' + str(lr))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### A2GradInc "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "A2GradInc\n",
      "Epoch 0, training loss 0.6602763776485593, time passed 0m 21s\n",
      "Validation loss 0.5054242114502093\n",
      "Epoch 2, training loss 0.4349289565605807, time passed 1m 2s\n",
      "Validation loss 0.4237436142879699\n",
      "Epoch 4, training loss 0.400614009719051, time passed 1m 48s\n",
      "Validation loss 0.3957431950432896\n"
     ]
    }
   ],
   "source": [
    "for lr in lrs:\n",
    "    model = LR()\n",
    "    print('A2GradInc')\n",
    "    optimizer = A2GradInc(model.parameters())\n",
    "    tr_loss['A2GradInc'][lr], times, val_loss['A2GradInc'][lr], lips, grad = train(model, trainloader, criterion, \n",
    "                                                                         optimizer, n_epochs=n_epochs, \n",
    "                                                                         print_every=print_every, \n",
    "                                                                         validloader=validloader)\n",
    "    states = {\n",
    "            'epoch': n_epochs,\n",
    "            'state_dict': model.state_dict(),\n",
    "            'optimizer': optimizer.state_dict(),\n",
    "            'tr_loss' : tr_loss['A2GradInc'][lr],\n",
    "            'val_loss' : val_loss['A2GradInc'][lr],\n",
    "            'lips' : lips,\n",
    "            'grad' : grad,\n",
    "            'times' : times\n",
    "             }\n",
    "    torch.save(states, './MNIST/LR_A2GradInc')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### A2GradExp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "A2GradExp  lr=0.001:\n",
      "Epoch 0, training loss 0.890304833246765, time passed 0m 16s\n",
      "Validation loss 0.6959082544174604\n",
      "Epoch 2, training loss 0.5943733410107238, time passed 1m 4s\n",
      "Validation loss 0.5783207188677183\n",
      "Epoch 4, training loss 0.5419563035308936, time passed 1m 45s\n",
      "Validation loss 0.5347908401390535\n"
     ]
    }
   ],
   "source": [
    "for lr in lrs:\n",
    "    model = LR()\n",
    "    print(\"A2GradExp  lr={}:\".format(lr))\n",
    "    optimizer = A2GradExp(model.parameters(), beta=10, lips=10, rho=0.9)\n",
    "    tr_loss['A2GradExp'][lr], times, val_loss['A2GradExp'][lr], lips, grad = train(model, trainloader, criterion, \n",
    "                                                                         optimizer, n_epochs=n_epochs, \n",
    "                                                                         print_every=print_every, \n",
    "                                                                         validloader=validloader)\n",
    "    states = {\n",
    "            'epoch': n_epochs,\n",
    "            'state_dict': model.state_dict(),\n",
    "            'optimizer': optimizer.state_dict(),\n",
    "            'tr_loss' : tr_loss['A2GradExp'][lr],\n",
    "            'val_loss' : val_loss['A2GradExp'][lr],\n",
    "            'lips' : lips,\n",
    "            'grad' : grad,\n",
    "            'times' : times\n",
    "             }\n",
    "    torch.save(states, './MNIST/LR_A2GradExp')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### A2GradUni"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "A2GradUni  lr=0.001:\n",
      "Epoch 0, training loss 0.8777040629732776, time passed 0m 21s\n",
      "Validation loss 0.6466419323036516\n",
      "Epoch 2, training loss 0.5362675025063032, time passed 1m 2s\n",
      "Validation loss 0.5200954246636614\n",
      "Epoch 4, training loss 0.4826502767078481, time passed 1m 52s\n",
      "Validation loss 0.47515102096064293\n"
     ]
    }
   ],
   "source": [
    "for lr in lrs:\n",
    "    model = LR()\n",
    "    print(\"A2GradUni  lr={}:\".format(lr))\n",
    "    optimizer = A2GradUni(model.parameters())\n",
    "    tr_loss['A2GradUni'][lr], times, val_loss['A2GradUni'][lr], lips, grad = train(model, trainloader, criterion, \n",
    "                                                                         optimizer, n_epochs=n_epochs, \n",
    "                                                                         print_every=print_every, \n",
    "                                                                         validloader=validloader)\n",
    "    states = {\n",
    "            'epoch': n_epochs,\n",
    "            'state_dict': model.state_dict(),\n",
    "            'optimizer': optimizer.state_dict(),\n",
    "            'tr_loss' : tr_loss['A2GradUni'][lr],\n",
    "            'val_loss' : val_loss['A2GradUni'][lr],\n",
    "            'lips' : lips,\n",
    "            'grad' : grad,\n",
    "            'times' : times\n",
    "             }\n",
    "    torch.save(states, './MNIST/LR_A2GradUni')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### AMSgrad"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "b1CbtL-F0xaF"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 0, training loss 0.4065497466627148, time passed 0m 19s\n",
      "Validation loss 0.42612661495290843\n",
      "Epoch 2, training loss 0.32533949986387606, time passed 1m 10s\n",
      "Validation loss 0.37711557346290964\n",
      "Epoch 4, training loss 0.31078923913632767, time passed 1m 55s\n",
      "Validation loss 0.3057402240929971\n"
     ]
    }
   ],
   "source": [
    "for lr in lrs:\n",
    "    model = LR()\n",
    "    optimizer = torch.optim.Adam(model.parameters(), lr=lr, betas=(0.9, 0.99), eps=1e-8, amsgrad=True)\n",
    "    tr_loss['amsgrad'], times, val_loss['amsgrad'], lips, grad = train(model, trainloader, criterion, optimizer, \n",
    "                                                           n_epochs=n_epochs, print_every=print_every, \n",
    "                                                           validloader=validloader)\n",
    "                                                         \n",
    "    states = {\n",
    "            'epoch': n_epochs,\n",
    "            'state_dict': model.state_dict(),\n",
    "            'optimizer': optimizer.state_dict(),\n",
    "            'tr_loss' : tr_loss['amsgrad'],\n",
    "            'val_loss' : val_loss['amsgrad'],\n",
    "            'lips' : lips,\n",
    "            'grad' : grad,\n",
    "            'times' : times\n",
    "         }\n",
    "    torch.save(states, './MNIST/LR_amsgrad')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Accelerated SGD"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 0, training loss 0.40764180013083695, time passed 0m 20s\n",
      "Validation loss 0.40028829447166864\n",
      "Epoch 2, training loss 0.30266382122565666, time passed 1m 7s\n",
      "Validation loss 0.36440590935486555\n",
      "Epoch 4, training loss 0.28831493548777254, time passed 1m 53s\n",
      "Validation loss 0.31914368522866027\n"
     ]
    }
   ],
   "source": [
    "model = LR()\n",
    "optimizer = AccSGD(model.parameters(), lr=0.001, kappa = 1000.0, xi = 10.0)\n",
    "tr_loss['accSGD'], times, val_loss['accSGD'], lips, grad = train(model, trainloader, criterion, optimizer, \n",
    "                                                           n_epochs=n_epochs, print_every=print_every, \n",
    "                                                           validloader=validloader)\n",
    "                                          \n",
    "states = {\n",
    "            'epoch': n_epochs,\n",
    "            'state_dict': model.state_dict(),\n",
    "            'optimizer': optimizer.state_dict(),\n",
    "            'tr_loss' : tr_loss['accSGD'],\n",
    "            'val_loss' : val_loss['accSGD'],\n",
    "            'lips' : lips,\n",
    "            'grad' : grad,\n",
    "            'times' : times\n",
    "         }\n",
    "torch.save(states, './MNIST/LR_accSGD')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Adaptive SGD"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 0, training loss 0.4637937442886873, time passed 0m 31s\n",
      "Validation loss 0.4307740188288825\n",
      "Epoch 2, training loss 0.38815918782290354, time passed 1m 32s\n",
      "Validation loss 0.43152904974208994\n",
      "Epoch 4, training loss 0.38372424873438565, time passed 2m 27s\n",
      "Validation loss 0.4823033701936678\n"
     ]
    }
   ],
   "source": [
    "for lr in lrs:\n",
    "    model = LR()\n",
    "    optimizer = SUG(model.parameters(), l_0=lr, momentum=0, nesterov=False)\n",
    "    tr_loss['sug'], times, val_loss['sug'], lips, grad = train(model, trainloader, criterion, optimizer, \n",
    "                                                           n_epochs=n_epochs, print_every=print_every, \n",
    "                                                           validloader=validloader)\n",
    "    states = {\n",
    "            'epoch': n_epochs,\n",
    "            'state_dict': model.state_dict(),\n",
    "            'optimizer': optimizer.state_dict(),\n",
    "            'tr_loss' : tr_loss['sug'],\n",
    "            'val_loss' : val_loss['sug'],\n",
    "            'lips' : lips,\n",
    "            'grad' : grad,\n",
    "            'times' : times\n",
    "         }\n",
    "    torch.save(states, './MNIST/LR_sug')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "uJ24b8FFpIQU"
   },
   "source": [
    "### FC neural network"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_epochs = 8"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Adam"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 374
    },
    "colab_type": "code",
    "id": "KlkSl5x-vPRW",
    "outputId": "0c739f08-c0ac-4326-b69d-870d51daa4a1"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Adam  lr=0.001:\n",
      "Epoch 0, training loss 0.31472805445884217, time passed 2m 17s\n",
      "Validation loss 0.33562005795256833\n",
      "Epoch 2, training loss 0.1640293464746599, time passed 7m 56s\n",
      "Validation loss 0.16987620446443386\n",
      "Epoch 4, training loss 0.1375037254112192, time passed 13m 44s\n",
      "Validation loss 0.1692539849275083\n",
      "Epoch 6, training loss 0.12588283565029146, time passed 20m 30s\n",
      "Validation loss 0.17024462949287886\n"
     ]
    }
   ],
   "source": [
    "for lr in lrs:\n",
    "    model = FC()\n",
    "    print(\"Adam  lr={}:\".format(lr))\n",
    "    optimizer = optim.Adam(model.parameters(), lr=lr)\n",
    "    tr_loss['Adam'][lr], times, val_loss['Adam'][lr], lips, grad = train(model, trainloader, criterion, \n",
    "                                                                         optimizer, n_epochs=n_epochs, \n",
    "                                                                         print_every=print_every, \n",
    "                                                                         validloader=validloader)\n",
    "    states = {\n",
    "            'epoch': n_epochs,\n",
    "            'state_dict': model.state_dict(),\n",
    "            'optimizer': optimizer.state_dict(),\n",
    "            'tr_loss' : tr_loss['Adam'][lr],\n",
    "            'val_loss' : val_loss['Adam'][lr],\n",
    "            'lips' : lips,\n",
    "            'grad' : grad,\n",
    "            'times' : times\n",
    "             }\n",
    "    torch.save(states, './MNIST/FC_Adam_' + str(lr))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### AMSgrad"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 119
    },
    "colab_type": "code",
    "id": "ODga9MkTvnRl",
    "outputId": "e1949814-6533-4078-9b5e-cf193b98398f"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 0, training loss 0.2954024291499492, time passed 2m 8s\n",
      "Validation loss 0.19679289057075633\n",
      "Epoch 2, training loss 0.11406174258036991, time passed 7m 39s\n",
      "Validation loss 0.14818219451598985\n",
      "Epoch 4, training loss 0.07966713166552923, time passed 13m 13s\n",
      "Validation loss 0.10309265803739515\n",
      "Epoch 6, training loss 0.06052994809905093, time passed 18m 52s\n",
      "Validation loss 0.10621145213906528\n"
     ]
    }
   ],
   "source": [
    "for lr in lrs:\n",
    "    model = FC()\n",
    "    optimizer = torch.optim.Adam(model.parameters(), lr=lr, betas=(0.9, 0.99), eps=1e-8, amsgrad=True)\n",
    "    tr_loss['amsgrad'], times, val_loss['amsgrad'], lips, grad = train(model, trainloader, criterion, optimizer, \n",
    "                                                           n_epochs=n_epochs, print_every=print_every, \n",
    "                                                           validloader=validloader)\n",
    "                                                         \n",
    "    states = {\n",
    "            'epoch': n_epochs,\n",
    "            'state_dict': model.state_dict(),\n",
    "            'optimizer': optimizer.state_dict(),\n",
    "            'tr_loss' : tr_loss['amsgrad'],\n",
    "            'val_loss' : val_loss['amsgrad'],\n",
    "            'lips' : lips,\n",
    "            'grad' : grad,\n",
    "            'times' : times\n",
    "         }\n",
    "    torch.save(states, './MNIST/FC_amsgrad')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Accelerated SGD"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "H8y30s5VE4px"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 0, training loss 0.4004168582505848, time passed 0m 36s\n",
      "Validation loss 0.26797187960428787\n",
      "Epoch 2, training loss 0.15343301741924836, time passed 1m 55s\n",
      "Validation loss 0.17337015810517428\n",
      "Epoch 4, training loss 0.10210626035243803, time passed 3m 13s\n",
      "Validation loss 0.1559533597549474\n",
      "Epoch 6, training loss 0.07533009543270398, time passed 4m 37s\n",
      "Validation loss 0.10859271298254157\n"
     ]
    }
   ],
   "source": [
    "model = FC()\n",
    "optimizer = AccSGD(model.parameters(), lr=0.001, kappa = 1000.0, xi = 10.0)\n",
    "tr_loss['accSGD'], times, val_loss['accSGD'], lips, grad = train(model, trainloader, criterion, optimizer, \n",
    "                                                           n_epochs=n_epochs, print_every=print_every, \n",
    "                                                           validloader=validloader)\n",
    "                                          \n",
    "states = {\n",
    "            'epoch': n_epochs,\n",
    "            'state_dict': model.state_dict(),\n",
    "            'optimizer': optimizer.state_dict(),\n",
    "            'tr_loss' : tr_loss['accSGD'],\n",
    "            'val_loss' : val_loss['accSGD'],\n",
    "            'lips' : lips,\n",
    "            'grad' : grad,\n",
    "            'times' : times\n",
    "         }\n",
    "torch.save(states, './MNIST/FC_accSGD')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Adaptive SGD"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 0, training loss 0.31520323795340593, time passed 1m 18s\n",
      "Validation loss 0.17935882584950086\n",
      "Epoch 2, training loss 0.11782419556486382, time passed 4m 13s\n",
      "Validation loss 0.11791805404194093\n",
      "Epoch 4, training loss 0.08247277856162973, time passed 7m 37s\n",
      "Validation loss 0.11297670972776604\n",
      "Epoch 6, training loss 0.06494110676108814, time passed 10m 53s\n",
      "Validation loss 0.11238901866388724\n"
     ]
    }
   ],
   "source": [
    "for lr in lrs:\n",
    "    model = FC()\n",
    "    optimizer = SUG(model.parameters(), l_0=lr, momentum=0, nesterov=False)\n",
    "    tr_loss['sug'], times, val_loss['sug'], lips, grad = train(model, trainloader, criterion, optimizer, \n",
    "                                                           n_epochs=n_epochs, print_every=print_every, \n",
    "                                                           validloader=validloader)\n",
    "    states = {\n",
    "            'epoch': n_epochs,\n",
    "            'state_dict': model.state_dict(),\n",
    "            'optimizer': optimizer.state_dict(),\n",
    "            'tr_loss' : tr_loss['sug'],\n",
    "            'val_loss' : val_loss['sug'],\n",
    "            'lips' : lips,\n",
    "            'grad' : grad,\n",
    "            'times' : times\n",
    "         }\n",
    "    torch.save(states, './MNIST/FC_sug')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### A2GradInc "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "A2GradInc\n",
      "Epoch 0, training loss 0.3405313328328119, time passed 0m 50s\n",
      "Validation loss 0.40539793937192486\n",
      "Epoch 2, training loss 0.2502931321926777, time passed 2m 36s\n",
      "Validation loss 0.25179623060090744\n",
      "Epoch 4, training loss 0.19747851773686237, time passed 4m 52s\n",
      "Validation loss 0.2213085105865465\n",
      "Epoch 6, training loss 0.19407693164069398, time passed 6m 41s\n",
      "Validation loss 0.20182552583856866\n"
     ]
    }
   ],
   "source": [
    "for lr in lrs:\n",
    "    model = FC()\n",
    "    print('A2GradInc')\n",
    "    optimizer = A2GradInc(model.parameters(), beta=1, lips=10)\n",
    "    tr_loss['A2GradInc'][lr], times, val_loss['A2GradInc'][lr], lips, grad = train(model, trainloader, criterion, \n",
    "                                                                         optimizer, n_epochs=n_epochs, \n",
    "                                                                         print_every=print_every, \n",
    "                                                                         validloader=validloader)\n",
    "    states = {\n",
    "            'epoch': n_epochs,\n",
    "            'state_dict': model.state_dict(),\n",
    "            'optimizer': optimizer.state_dict(),\n",
    "            'tr_loss' : tr_loss['A2GradInc'][lr],\n",
    "            'val_loss' : val_loss['A2GradInc'][lr],\n",
    "            'lips' : lips,\n",
    "            'grad' : grad,\n",
    "            'times' : times\n",
    "             }\n",
    "    torch.save(states, './MNIST/FC_A2GradInc')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### A2GradUni"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "A2GradUni  lr=0.001:\n",
      "Epoch 0, training loss 0.3645090667142411, time passed 0m 49s\n",
      "Validation loss 0.3461166744949872\n",
      "Epoch 2, training loss 0.2715470625102375, time passed 2m 48s\n",
      "Validation loss 0.25930503801856375\n",
      "Epoch 4, training loss 0.21909287481441134, time passed 4m 37s\n",
      "Validation loss 0.20832079588015273\n",
      "Epoch 6, training loss 0.1821865990205095, time passed 6m 28s\n",
      "Validation loss 0.17309671778210325\n"
     ]
    }
   ],
   "source": [
    "for lr in lrs:\n",
    "    model = FC()\n",
    "    print(\"A2GradUni  lr={}:\".format(lr))\n",
    "    optimizer = A2GradUni(model.parameters(), beta=1, lips=10)\n",
    "    tr_loss['A2GradUni'][lr], times, val_loss['A2GradUni'][lr], lips, grad = train(model, trainloader, criterion, \n",
    "                                                                         optimizer, n_epochs=n_epochs, \n",
    "                                                                         print_every=print_every, \n",
    "                                                                         validloader=validloader)\n",
    "    states = {\n",
    "            'epoch': n_epochs,\n",
    "            'state_dict': model.state_dict(),\n",
    "            'optimizer': optimizer.state_dict(),\n",
    "            'tr_loss' : tr_loss['A2GradUni'][lr],\n",
    "            'val_loss' : val_loss['A2GradUni'][lr],\n",
    "            'lips' : lips,\n",
    "            'grad' : grad,\n",
    "            'times' : times\n",
    "             }\n",
    "    torch.save(states, './MNIST/FC_A2GradUni')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### A2GradExp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "A2GradExp  lr=0.001:\n",
      "Epoch 0, training loss 0.38943213840140223, time passed 0m 48s\n",
      "Validation loss 0.3454775995363284\n",
      "Epoch 2, training loss 0.289244648559167, time passed 2m 49s\n",
      "Validation loss 0.28307165995922656\n",
      "Epoch 4, training loss 0.2758141870881653, time passed 4m 34s\n",
      "Validation loss 0.2541869924883675\n",
      "Epoch 6, training loss 0.2549206986760371, time passed 6m 14s\n",
      "Validation loss 0.2627989425029475\n"
     ]
    }
   ],
   "source": [
    "for lr in lrs:\n",
    "    model = FC()\n",
    "    print(\"A2GradExp  lr={}:\".format(lr))\n",
    "    optimizer = A2GradExp(model.parameters(), beta=1, lips=5, rho=0.99)\n",
    "    tr_loss['A2GradExp'][lr], times, val_loss['A2GradExp'][lr], lips, grad = train(model, trainloader, criterion, \n",
    "                                                                         optimizer, n_epochs=n_epochs, \n",
    "                                                                         print_every=print_every, \n",
    "                                                                         validloader=validloader)\n",
    "    states = {\n",
    "            'epoch': n_epochs,\n",
    "            'state_dict': model.state_dict(),\n",
    "            'optimizer': optimizer.state_dict(),\n",
    "            'tr_loss' : tr_loss['A2GradExp'][lr],\n",
    "            'val_loss' : val_loss['A2GradExp'][lr],\n",
    "            'lips' : lips,\n",
    "            'grad' : grad,\n",
    "            'times' : times\n",
    "             }\n",
    "    torch.save(states, './MNIST/FC_A2GradExp')"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "include_colab_link": true,
   "name": "MNIST.ipynb",
   "provenance": [],
   "toc_visible": true,
   "version": "0.3.2"
  },
  "kernelspec": {
   "display_name": "Python 3",
   "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.7.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}