NPCodeRegression_z_mu.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.utils.data\n",
    "from torch import nn, optim\n",
    "from torch.nn import functional as F\n",
    "from torchvision import datasets, transforms\n",
    "from torchvision.utils import save_image\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from sklearn.gaussian_process.kernels import RBF\n",
    "import os\n",
    "import random"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size=512\n",
    "epochs=300\n",
    "seed=1\n",
    "cuda=False and torch.cuda.is_available()\n",
    "log_interval=1\n",
    "r_dim=128\n",
    "z_dim=128\n",
    "result_path=\"results_np_z_y_hat_parmu_reg/\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generateData(num_pt=50,length_scale=1.0,plot=False):\n",
    "    #np.random.seed(2)\n",
    "    x=np.random.uniform(-3,3,num_pt)[:,None]\n",
    "    x.sort(axis=0)\n",
    "    k=RBF(length_scale=length_scale)(x,x)\n",
    "    y=np.random.multivariate_normal(np.zeros(num_pt),k)[:,None]\n",
    "    N=np.random.randint(1, num_pt)\n",
    "    loc=np.random.choice(num_pt,N)\n",
    "    x_context=x[loc]\n",
    "    y_context=y[loc]\n",
    "    if plot:\n",
    "        plt.plot(x_context,y_context,'ko')\n",
    "        plt.plot(x,y)\n",
    "    return x_context,y_context,x,y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "class NP(nn.Module):\n",
    "    def __init__(self, r_dim,z_dim):\n",
    "        super(NP, self).__init__()\n",
    "        self.r_dim = r_dim\n",
    "        self.z_dim = z_dim\n",
    "    \n",
    "        self.h_1 = nn.Linear(3, 400)\n",
    "        self.h_2 = nn.Linear(400, 400)\n",
    "        self.h_3 = nn.Linear(400, self.r_dim)\n",
    "\n",
    "        self.r_to_z_mean = nn.Linear(self.r_dim, self.z_dim)\n",
    "        self.r_to_z_logvar = nn.Linear(self.r_dim, self.z_dim)\n",
    "\n",
    "        self.g_1 = nn.Linear(self.z_dim + 2, 400)\n",
    "        self.g_2 = nn.Linear(400, 400)\n",
    "        self.g_3 = nn.Linear(400, 400)\n",
    "        self.g_4 = nn.Linear(400, 400)\n",
    "        self.g_y_mu = nn.Linear(400, 1)\n",
    "        self.g_y_sigma = nn.Linear(400, 1)\n",
    "        \n",
    "\n",
    "    def h(self, x_y):\n",
    "        x_y = F.relu(self.h_1(x_y))\n",
    "        x_y = F.relu(self.h_2(x_y))\n",
    "        x_y = F.relu(self.h_3(x_y))\n",
    "        return x_y\n",
    "\n",
    "    def aggregate(self, r):\n",
    "        return torch.mean(r, dim=1)\n",
    "\n",
    "    def g(self,z_sample, x_target):\n",
    "        z_et_x = torch.cat([z_sample, x_target], dim=2)\n",
    "        input = F.relu(self.g_1(z_et_x))\n",
    "        input = F.relu(self.g_2(input))\n",
    "        input = F.relu(self.g_3(input))\n",
    "        input = F.relu(self.g_4(input))\n",
    "        y_mu=self.g_y_mu(input)\n",
    "        y_sigma=sigma = 0.1 + 0.9 * F.softplus(self.g_y_sigma(input))\n",
    "        return y_mu,y_sigma\n",
    "    \n",
    "    \n",
    "    def xy_to_z_params(self, x, y):\n",
    "        \n",
    "        x_y = torch.cat([x, y], dim=2)\n",
    "        \n",
    "        r_i = self.h(x_y)\n",
    "        r = self.aggregate(r_i)\n",
    "        mu = self.r_to_z_mean(r)\n",
    "        logvar = self.r_to_z_logvar(r)\n",
    "        sigma=0.1+0.9*torch.sigmoid(logvar)\n",
    "        return mu, sigma\n",
    "\n",
    "    def forward(self, x_context, y_context, x_all=None, y_all=None):\n",
    "        \n",
    "        #produire z\n",
    "        z_context_mu,z_context_sigma = self.xy_to_z_params(x_context, y_context)  # (mu, logvar) of z\n",
    "        q_context = Normal(z_context_mu, z_context_sigma)\n",
    "        # reconstruct the whole image including the provided context points\n",
    "        x_target = x_grid.expand(y_context.shape[0], -1, -1)\n",
    "        \n",
    "        if self.training:  # loss function will try to keep z_context close to z_all         \n",
    "            z_target_mu,z_target_sigma = self.xy_to_z_params(x_context, y_context)\n",
    "            q_target = Normal(z_target_mu, z_target_sigma) \n",
    "            z_sample = q_target.rsample()\n",
    "            \n",
    "            z_sample = z_sample.unsqueeze(1).expand(-1, 784, -1)\n",
    "            \n",
    "            # Get parameters of output distribution\n",
    "            y_pred_mu, y_pred_sigma = self.g(z_sample,x_target)\n",
    "            p_y_pred = Normal(y_pred_mu, y_pred_sigma)\n",
    "\n",
    "            return p_y_pred, q_target, q_context\n",
    "        else:  # at test time we don't have the image so we use only the context\n",
    "            z_sample = q_context.rsample()\n",
    "            z_sample = z_sample.unsqueeze(1).expand(-1, 784, -1)\n",
    "            # Predict target points based on context\n",
    "            y_pred_mu, y_pred_sigma = self.g(z_sample,x_target)\n",
    "            p_y_pred = Normal(y_pred_mu, y_pred_sigma)\n",
    "            return p_y_pred,q_context,q_context"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torch.distributions import Normal\n",
    "from torch.distributions.kl import kl_divergence\n",
    "\n",
    "def np_loss(p_y_pred, y_target, q_target, q_context):\n",
    "    \n",
    "    #return logprob + KLD\n",
    "    log_likelihood = p_y_pred.log_prob(y_target).mean(dim=0).sum()\n",
    "    # KL has shape (batch_size, r_dim). Take mean over batch and sum over\n",
    "    # r_dim (since r_dim is dimension of normal distribution)\n",
    "    kl = kl_divergence(q_target, q_context).mean(dim=0).sum()\n",
    "    return -log_likelihood + kl\n",
    "    \n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = NP(r_dim,z_dim).to(device)\n",
    "optimizer = optim.Adam(model.parameters(), lr=5e-4)\n",
    "x_grid = generate_grid(28, 28)\n",
    "os.makedirs(result_path, exist_ok=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(epoch):\n",
    "    model.train()\n",
    "    train_loss = 0\n",
    "    for batch_idx, (y_all, _) in enumerate(train_loader):\n",
    "        \n",
    "        batch_size = y_all.shape[0]\n",
    "        \n",
    "        y_all = y_all.to(device).view(batch_size, -1, 1)\n",
    "        \n",
    "        N = random.randint(1, 784)  # number of context points\n",
    "        context_idx = get_context_idx(N)\n",
    "        x_context = idx_to_x(context_idx, batch_size)\n",
    "        y_context = idx_to_y(context_idx, y_all)\n",
    "        x_all = x_grid.expand(batch_size, -1, -1)\n",
    "\n",
    "        optimizer.zero_grad()\n",
    "        \n",
    "        p_y_pred, q_target, q_context = model(x_context, y_context, x_all, y_all)\n",
    "        \n",
    "        loss = np_loss(p_y_pred, y_all, q_target, q_context)\n",
    "        loss.backward()\n",
    "        train_loss += loss.item()\n",
    "        optimizer.step()\n",
    "        if batch_idx % log_interval == 0:\n",
    "            print('Train Epoch: {} [{}/{} ({:.0f}%)]\\tLoss: {:.6f}'.format(\n",
    "                epoch, batch_idx * len(y_all), len(train_loader.dataset),\n",
    "                       100. * batch_idx / len(train_loader),\n",
    "                       loss.item() / len(y_all)))\n",
    "    print('====> Epoch: {} Average loss: {:.4f}'.format(\n",
    "        epoch, train_loss / len(train_loader.dataset)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "def test(epoch):\n",
    "    model.eval()\n",
    "    test_loss = 0\n",
    "    with torch.no_grad():\n",
    "        for i, (y_all, _) in enumerate(test_loader):\n",
    "            y_all = y_all.to(device).view(y_all.shape[0], -1, 1)\n",
    "            batch_size = y_all.shape[0]\n",
    "\n",
    "            N = 300\n",
    "            context_idx = get_context_idx(N)\n",
    "            x_context = idx_to_x(context_idx, batch_size)\n",
    "            y_context = idx_to_y(context_idx, y_all)\n",
    "            \n",
    "            p_y_pred, q_target, q_context = model(x_context, y_context)\n",
    "        \n",
    "            \n",
    "            test_loss += np_loss(p_y_pred, y_all, q_target, q_context).item()\n",
    "\n",
    "            if i == 0:  # save PNG of reconstructed examples\n",
    "                plot_Ns = [10, 100, 300, 784]\n",
    "                num_examples = min(batch_size, 16)\n",
    "                for N in plot_Ns:\n",
    "                    recons = []\n",
    "                    context_idx = get_context_idx(N)\n",
    "                    x_context = idx_to_x(context_idx, batch_size)\n",
    "                    y_context = idx_to_y(context_idx, y_all)\n",
    "                    for d in range(5):\n",
    "                        p_y_pred, _, _ = model(x_context, y_context)\n",
    "                        \n",
    "                        recons.append(p_y_pred.rsample()[:num_examples])     \n",
    "                    recons = torch.cat(recons).view(-1, 1, 28, 28).expand(-1, 3, -1, -1)\n",
    "                    background = torch.tensor([0., 0., 1.], device=device)\n",
    "                    background = background.view(1, -1, 1).expand(num_examples, 3, 784).contiguous()\n",
    "                    context_pixels = y_all[:num_examples].view(num_examples, 1, -1)[:, :, context_idx]\n",
    "                    context_pixels = context_pixels.expand(num_examples, 3, -1)\n",
    "                    background[:, :, context_idx] = context_pixels\n",
    "                    comparison = torch.cat([background.view(-1, 3, 28, 28),\n",
    "                                            recons])\n",
    "                    save_image(comparison.cpu(),\n",
    "                               result_path+\"ep_\" + str(epoch) +\n",
    "                               \"_nps_\" + str(N) + \".png\", nrow=num_examples)\n",
    "\n",
    "    test_loss /= len(test_loader.dataset)\n",
    "    print('====> Test set loss: {:.4f}'.format(test_loss))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train Epoch: 14 [0/60000 (0%)]\tLoss: -0.767608\n",
      "Train Epoch: 14 [512/60000 (1%)]\tLoss: -0.751735\n",
      "Train Epoch: 14 [1024/60000 (2%)]\tLoss: -0.757991\n",
      "Train Epoch: 14 [1536/60000 (3%)]\tLoss: -0.745510\n",
      "Train Epoch: 14 [2048/60000 (3%)]\tLoss: -0.759890\n",
      "Train Epoch: 14 [2560/60000 (4%)]\tLoss: -0.755115\n",
      "Train Epoch: 14 [3072/60000 (5%)]\tLoss: -0.758321\n",
      "Train Epoch: 14 [3584/60000 (6%)]\tLoss: -0.785377\n",
      "Train Epoch: 14 [4096/60000 (7%)]\tLoss: -0.796116\n",
      "Train Epoch: 14 [4608/60000 (8%)]\tLoss: -0.762706\n",
      "Train Epoch: 14 [5120/60000 (8%)]\tLoss: -0.769937\n",
      "Train Epoch: 14 [5632/60000 (9%)]\tLoss: -0.725292\n",
      "Train Epoch: 14 [6144/60000 (10%)]\tLoss: -0.764569\n",
      "Train Epoch: 14 [6656/60000 (11%)]\tLoss: -0.746527\n",
      "Train Epoch: 14 [7168/60000 (12%)]\tLoss: -0.748802\n",
      "Train Epoch: 14 [7680/60000 (13%)]\tLoss: -0.766164\n",
      "Train Epoch: 14 [8192/60000 (14%)]\tLoss: -0.738145\n",
      "Train Epoch: 14 [8704/60000 (14%)]\tLoss: -0.738172\n",
      "Train Epoch: 14 [9216/60000 (15%)]\tLoss: -0.753874\n",
      "Train Epoch: 14 [9728/60000 (16%)]\tLoss: -0.769690\n",
      "Train Epoch: 14 [10240/60000 (17%)]\tLoss: -0.753884\n",
      "Train Epoch: 14 [10752/60000 (18%)]\tLoss: -0.762737\n",
      "Train Epoch: 14 [11264/60000 (19%)]\tLoss: -0.761026\n",
      "Train Epoch: 14 [11776/60000 (19%)]\tLoss: -0.781266\n",
      "Train Epoch: 14 [12288/60000 (20%)]\tLoss: -0.645938\n",
      "Train Epoch: 14 [12800/60000 (21%)]\tLoss: -0.744303\n",
      "Train Epoch: 14 [13312/60000 (22%)]\tLoss: -0.709032\n",
      "Train Epoch: 14 [13824/60000 (23%)]\tLoss: -0.735938\n",
      "Train Epoch: 14 [14336/60000 (24%)]\tLoss: -0.753478\n",
      "Train Epoch: 14 [14848/60000 (25%)]\tLoss: -0.766261\n",
      "Train Epoch: 14 [15360/60000 (25%)]\tLoss: -0.745685\n",
      "Train Epoch: 14 [15872/60000 (26%)]\tLoss: -0.762205\n",
      "Train Epoch: 14 [16384/60000 (27%)]\tLoss: -0.762588\n",
      "Train Epoch: 14 [16896/60000 (28%)]\tLoss: -0.745737\n",
      "Train Epoch: 14 [17408/60000 (29%)]\tLoss: -0.719646\n",
      "Train Epoch: 14 [17920/60000 (30%)]\tLoss: -0.760934\n",
      "Train Epoch: 14 [18432/60000 (31%)]\tLoss: -0.746924\n",
      "Train Epoch: 14 [18944/60000 (31%)]\tLoss: -0.753421\n",
      "Train Epoch: 14 [19456/60000 (32%)]\tLoss: -0.764261\n",
      "Train Epoch: 14 [19968/60000 (33%)]\tLoss: -0.770431\n",
      "Train Epoch: 14 [20480/60000 (34%)]\tLoss: -0.763425\n",
      "Train Epoch: 14 [20992/60000 (35%)]\tLoss: -0.748565\n",
      "Train Epoch: 14 [21504/60000 (36%)]\tLoss: -0.768681\n",
      "Train Epoch: 14 [22016/60000 (36%)]\tLoss: -0.757309\n",
      "Train Epoch: 14 [22528/60000 (37%)]\tLoss: -0.746773\n",
      "Train Epoch: 14 [23040/60000 (38%)]\tLoss: -0.754103\n",
      "Train Epoch: 14 [23552/60000 (39%)]\tLoss: -0.787085\n",
      "Train Epoch: 14 [24064/60000 (40%)]\tLoss: -0.760307\n",
      "Train Epoch: 14 [24576/60000 (41%)]\tLoss: -0.770748\n",
      "Train Epoch: 14 [25088/60000 (42%)]\tLoss: -0.768821\n",
      "Train Epoch: 14 [25600/60000 (42%)]\tLoss: -0.786627\n",
      "Train Epoch: 14 [26112/60000 (43%)]\tLoss: -0.782711\n",
      "Train Epoch: 14 [26624/60000 (44%)]\tLoss: -0.778872\n",
      "Train Epoch: 14 [27136/60000 (45%)]\tLoss: -0.782551\n",
      "Train Epoch: 14 [27648/60000 (46%)]\tLoss: -0.731375\n",
      "Train Epoch: 14 [28160/60000 (47%)]\tLoss: -0.770583\n",
      "Train Epoch: 14 [28672/60000 (47%)]\tLoss: -0.761174\n",
      "Train Epoch: 14 [29184/60000 (48%)]\tLoss: -0.755507\n",
      "Train Epoch: 14 [29696/60000 (49%)]\tLoss: -0.754082\n",
      "Train Epoch: 14 [30208/60000 (50%)]\tLoss: -0.784538\n",
      "Train Epoch: 14 [30720/60000 (51%)]\tLoss: -0.762835\n",
      "Train Epoch: 14 [31232/60000 (52%)]\tLoss: -0.788756\n",
      "Train Epoch: 14 [31744/60000 (53%)]\tLoss: -0.753092\n",
      "Train Epoch: 14 [32256/60000 (53%)]\tLoss: -0.770416\n",
      "Train Epoch: 14 [32768/60000 (54%)]\tLoss: -0.758359\n",
      "Train Epoch: 14 [33280/60000 (55%)]\tLoss: -0.759005\n",
      "Train Epoch: 14 [33792/60000 (56%)]\tLoss: -0.749253\n",
      "Train Epoch: 14 [34304/60000 (57%)]\tLoss: -0.745789\n",
      "Train Epoch: 14 [34816/60000 (58%)]\tLoss: -0.744360\n",
      "Train Epoch: 14 [35328/60000 (58%)]\tLoss: -0.781306\n",
      "Train Epoch: 14 [35840/60000 (59%)]\tLoss: -0.739453\n",
      "Train Epoch: 14 [36352/60000 (60%)]\tLoss: -0.751755\n",
      "Train Epoch: 14 [36864/60000 (61%)]\tLoss: -0.750809\n",
      "Train Epoch: 14 [37376/60000 (62%)]\tLoss: -0.767857\n",
      "Train Epoch: 14 [37888/60000 (63%)]\tLoss: -0.760995\n",
      "Train Epoch: 14 [38400/60000 (64%)]\tLoss: -0.789339\n",
      "Train Epoch: 14 [38912/60000 (64%)]\tLoss: -0.762213\n",
      "Train Epoch: 14 [39424/60000 (65%)]\tLoss: -0.764142\n",
      "Train Epoch: 14 [39936/60000 (66%)]\tLoss: -0.507518\n",
      "Train Epoch: 14 [40448/60000 (67%)]\tLoss: -0.714549\n",
      "Train Epoch: 14 [40960/60000 (68%)]\tLoss: -0.737011\n",
      "Train Epoch: 14 [41472/60000 (69%)]\tLoss: -0.762487\n",
      "Train Epoch: 14 [41984/60000 (69%)]\tLoss: -0.728341\n",
      "Train Epoch: 14 [42496/60000 (70%)]\tLoss: -0.776032\n",
      "Train Epoch: 14 [43008/60000 (71%)]\tLoss: -0.774549\n",
      "Train Epoch: 14 [43520/60000 (72%)]\tLoss: -0.757647\n",
      "Train Epoch: 14 [44032/60000 (73%)]\tLoss: -0.757184\n",
      "Train Epoch: 14 [44544/60000 (74%)]\tLoss: -0.735447\n",
      "Train Epoch: 14 [45056/60000 (75%)]\tLoss: -0.751760\n",
      "Train Epoch: 14 [45568/60000 (75%)]\tLoss: -0.751366\n",
      "Train Epoch: 14 [46080/60000 (76%)]\tLoss: -0.037937\n",
      "Train Epoch: 14 [46592/60000 (77%)]\tLoss: -0.727304\n",
      "Train Epoch: 14 [47104/60000 (78%)]\tLoss: -0.748431\n",
      "Train Epoch: 14 [47616/60000 (79%)]\tLoss: -0.759071\n",
      "Train Epoch: 14 [48128/60000 (80%)]\tLoss: -0.736939\n",
      "Train Epoch: 14 [48640/60000 (81%)]\tLoss: -0.722449\n",
      "Train Epoch: 14 [49152/60000 (81%)]\tLoss: -0.720153\n",
      "Train Epoch: 14 [49664/60000 (82%)]\tLoss: -0.747903\n",
      "Train Epoch: 14 [50176/60000 (83%)]\tLoss: -0.757241\n",
      "Train Epoch: 14 [50688/60000 (84%)]\tLoss: -0.774290\n",
      "Train Epoch: 14 [51200/60000 (85%)]\tLoss: -0.754579\n",
      "Train Epoch: 14 [51712/60000 (86%)]\tLoss: -0.721718\n",
      "Train Epoch: 14 [52224/60000 (86%)]\tLoss: -0.764431\n",
      "Train Epoch: 14 [52736/60000 (87%)]\tLoss: -0.760619\n",
      "Train Epoch: 14 [53248/60000 (88%)]\tLoss: -0.742988\n",
      "Train Epoch: 14 [53760/60000 (89%)]\tLoss: -0.753263\n",
      "Train Epoch: 14 [54272/60000 (90%)]\tLoss: -0.782225\n",
      "Train Epoch: 14 [54784/60000 (91%)]\tLoss: -0.749762\n",
      "Train Epoch: 14 [55296/60000 (92%)]\tLoss: -0.507746\n",
      "Train Epoch: 14 [55808/60000 (92%)]\tLoss: -0.759879\n",
      "Train Epoch: 14 [56320/60000 (93%)]\tLoss: -0.754430\n",
      "Train Epoch: 14 [56832/60000 (94%)]\tLoss: -0.753337\n",
      "Train Epoch: 14 [57344/60000 (95%)]\tLoss: -0.763436\n",
      "Train Epoch: 14 [57856/60000 (96%)]\tLoss: -0.746465\n",
      "Train Epoch: 14 [58368/60000 (97%)]\tLoss: -0.742771\n",
      "Train Epoch: 14 [58880/60000 (97%)]\tLoss: -0.749777\n",
      "Train Epoch: 14 [59392/60000 (98%)]\tLoss: -0.769111\n",
      "Train Epoch: 14 [11232/60000 (99%)]\tLoss: -4.190151\n",
      "====> Epoch: 14 Average loss: -0.7507\n",
      "====> Test set loss: -0.7794\n",
      "Train Epoch: 15 [0/60000 (0%)]\tLoss: -0.758506\n",
      "Train Epoch: 15 [512/60000 (1%)]\tLoss: -0.717395\n",
      "Train Epoch: 15 [1024/60000 (2%)]\tLoss: -0.760944\n",
      "Train Epoch: 15 [1536/60000 (3%)]\tLoss: -0.726471\n",
      "Train Epoch: 15 [2048/60000 (3%)]\tLoss: -0.757749\n",
      "Train Epoch: 15 [2560/60000 (4%)]\tLoss: -0.783231\n",
      "Train Epoch: 15 [3072/60000 (5%)]\tLoss: -0.769248\n",
      "Train Epoch: 15 [3584/60000 (6%)]\tLoss: -0.759533\n",
      "Train Epoch: 15 [4096/60000 (7%)]\tLoss: -0.755262\n",
      "Train Epoch: 15 [4608/60000 (8%)]\tLoss: -0.766732\n",
      "Train Epoch: 15 [5120/60000 (8%)]\tLoss: -0.774495\n",
      "Train Epoch: 15 [5632/60000 (9%)]\tLoss: -0.774726\n",
      "Train Epoch: 15 [6144/60000 (10%)]\tLoss: -0.759759\n",
      "Train Epoch: 15 [6656/60000 (11%)]\tLoss: -0.750593\n",
      "Train Epoch: 15 [7168/60000 (12%)]\tLoss: -0.768106\n",
      "Train Epoch: 15 [7680/60000 (13%)]\tLoss: -0.782044\n",
      "Train Epoch: 15 [8192/60000 (14%)]\tLoss: -0.749484\n",
      "Train Epoch: 15 [8704/60000 (14%)]\tLoss: -0.724378\n",
      "Train Epoch: 15 [9216/60000 (15%)]\tLoss: -0.783283\n",
      "Train Epoch: 15 [9728/60000 (16%)]\tLoss: -0.762435\n",
      "Train Epoch: 15 [10240/60000 (17%)]\tLoss: -0.733227\n",
      "Train Epoch: 15 [10752/60000 (18%)]\tLoss: -0.765873\n",
      "Train Epoch: 15 [11264/60000 (19%)]\tLoss: -0.748907\n",
      "Train Epoch: 15 [11776/60000 (19%)]\tLoss: -0.774797\n",
      "Train Epoch: 15 [12288/60000 (20%)]\tLoss: -0.765300\n",
      "Train Epoch: 15 [12800/60000 (21%)]\tLoss: -0.771478\n",
      "Train Epoch: 15 [13312/60000 (22%)]\tLoss: -0.776043\n",
      "Train Epoch: 15 [13824/60000 (23%)]\tLoss: -0.739413\n",
      "Train Epoch: 15 [14336/60000 (24%)]\tLoss: -0.775024\n",
      "Train Epoch: 15 [14848/60000 (25%)]\tLoss: -0.733231\n",
      "Train Epoch: 15 [15360/60000 (25%)]\tLoss: -0.768412\n",
      "Train Epoch: 15 [15872/60000 (26%)]\tLoss: -0.785151\n",
      "Train Epoch: 15 [16384/60000 (27%)]\tLoss: -0.793131\n",
      "Train Epoch: 15 [16896/60000 (28%)]\tLoss: -0.740503\n",
      "Train Epoch: 15 [17408/60000 (29%)]\tLoss: -0.748398\n",
      "Train Epoch: 15 [17920/60000 (30%)]\tLoss: -0.759415\n",
      "Train Epoch: 15 [18432/60000 (31%)]\tLoss: -0.771991\n",
      "Train Epoch: 15 [18944/60000 (31%)]\tLoss: -0.774925\n",
      "Train Epoch: 15 [19456/60000 (32%)]\tLoss: -0.772875\n",
      "Train Epoch: 15 [19968/60000 (33%)]\tLoss: -0.730482\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train Epoch: 15 [20480/60000 (34%)]\tLoss: -0.782769\n",
      "Train Epoch: 15 [20992/60000 (35%)]\tLoss: -0.712952\n",
      "Train Epoch: 15 [21504/60000 (36%)]\tLoss: -0.757199\n",
      "Train Epoch: 15 [22016/60000 (36%)]\tLoss: -0.750142\n",
      "Train Epoch: 15 [22528/60000 (37%)]\tLoss: -0.757212\n",
      "Train Epoch: 15 [23040/60000 (38%)]\tLoss: -0.792166\n",
      "Train Epoch: 15 [23552/60000 (39%)]\tLoss: -0.760809\n",
      "Train Epoch: 15 [24064/60000 (40%)]\tLoss: -0.764642\n",
      "Train Epoch: 15 [24576/60000 (41%)]\tLoss: -0.773502\n",
      "Train Epoch: 15 [25088/60000 (42%)]\tLoss: -0.764615\n",
      "Train Epoch: 15 [25600/60000 (42%)]\tLoss: -0.749009\n",
      "Train Epoch: 15 [26112/60000 (43%)]\tLoss: -0.785451\n",
      "Train Epoch: 15 [26624/60000 (44%)]\tLoss: -0.754639\n",
      "Train Epoch: 15 [27136/60000 (45%)]\tLoss: -0.742397\n",
      "Train Epoch: 15 [27648/60000 (46%)]\tLoss: -0.777659\n",
      "Train Epoch: 15 [28160/60000 (47%)]\tLoss: -0.771843\n",
      "Train Epoch: 15 [28672/60000 (47%)]\tLoss: -0.764871\n",
      "Train Epoch: 15 [29184/60000 (48%)]\tLoss: -0.623917\n",
      "Train Epoch: 15 [29696/60000 (49%)]\tLoss: -0.779492\n",
      "Train Epoch: 15 [30208/60000 (50%)]\tLoss: -0.769262\n",
      "Train Epoch: 15 [30720/60000 (51%)]\tLoss: -0.763533\n",
      "Train Epoch: 15 [31232/60000 (52%)]\tLoss: -0.769405\n",
      "Train Epoch: 15 [31744/60000 (53%)]\tLoss: -0.763176\n",
      "Train Epoch: 15 [32256/60000 (53%)]\tLoss: -0.772607\n",
      "Train Epoch: 15 [32768/60000 (54%)]\tLoss: -0.766494\n",
      "Train Epoch: 15 [33280/60000 (55%)]\tLoss: -0.746674\n",
      "Train Epoch: 15 [33792/60000 (56%)]\tLoss: -0.749112\n",
      "Train Epoch: 15 [34304/60000 (57%)]\tLoss: -0.743316\n",
      "Train Epoch: 15 [34816/60000 (58%)]\tLoss: -0.759168\n",
      "Train Epoch: 15 [35328/60000 (58%)]\tLoss: -0.708110\n",
      "Train Epoch: 15 [35840/60000 (59%)]\tLoss: -0.215858\n",
      "Train Epoch: 15 [36352/60000 (60%)]\tLoss: -0.769473\n",
      "Train Epoch: 15 [36864/60000 (61%)]\tLoss: -0.741134\n",
      "Train Epoch: 15 [37376/60000 (62%)]\tLoss: -0.515925\n",
      "Train Epoch: 15 [37888/60000 (63%)]\tLoss: -0.765758\n",
      "Train Epoch: 15 [38400/60000 (64%)]\tLoss: -0.726600\n",
      "Train Epoch: 15 [38912/60000 (64%)]\tLoss: -0.714401\n",
      "Train Epoch: 15 [39424/60000 (65%)]\tLoss: -0.705386\n",
      "Train Epoch: 15 [39936/60000 (66%)]\tLoss: -0.750393\n",
      "Train Epoch: 15 [40448/60000 (67%)]\tLoss: -0.747976\n",
      "Train Epoch: 15 [40960/60000 (68%)]\tLoss: -0.747548\n",
      "Train Epoch: 15 [41472/60000 (69%)]\tLoss: -0.736825\n",
      "Train Epoch: 15 [41984/60000 (69%)]\tLoss: -0.734258\n",
      "Train Epoch: 15 [42496/60000 (70%)]\tLoss: -0.729301\n",
      "Train Epoch: 15 [43008/60000 (71%)]\tLoss: -0.775928\n",
      "Train Epoch: 15 [43520/60000 (72%)]\tLoss: -0.738273\n",
      "Train Epoch: 15 [44032/60000 (73%)]\tLoss: -0.762835\n",
      "Train Epoch: 15 [44544/60000 (74%)]\tLoss: -0.530180\n",
      "Train Epoch: 15 [45056/60000 (75%)]\tLoss: -0.753887\n",
      "Train Epoch: 15 [45568/60000 (75%)]\tLoss: -0.743097\n",
      "Train Epoch: 15 [46080/60000 (76%)]\tLoss: -0.727156\n",
      "Train Epoch: 15 [46592/60000 (77%)]\tLoss: -0.726874\n",
      "Train Epoch: 15 [47104/60000 (78%)]\tLoss: -0.746382\n",
      "Train Epoch: 15 [47616/60000 (79%)]\tLoss: -0.740949\n",
      "Train Epoch: 15 [48128/60000 (80%)]\tLoss: -0.761707\n",
      "Train Epoch: 15 [48640/60000 (81%)]\tLoss: -0.678974\n",
      "Train Epoch: 15 [49152/60000 (81%)]\tLoss: -0.757794\n",
      "Train Epoch: 15 [49664/60000 (82%)]\tLoss: -0.749135\n",
      "Train Epoch: 15 [50176/60000 (83%)]\tLoss: -0.728451\n",
      "Train Epoch: 15 [50688/60000 (84%)]\tLoss: -0.721030\n",
      "Train Epoch: 15 [51200/60000 (85%)]\tLoss: -0.743011\n",
      "Train Epoch: 15 [51712/60000 (86%)]\tLoss: -0.754271\n",
      "Train Epoch: 15 [52224/60000 (86%)]\tLoss: -0.767036\n",
      "Train Epoch: 15 [52736/60000 (87%)]\tLoss: -0.749153\n",
      "Train Epoch: 15 [53248/60000 (88%)]\tLoss: -0.747792\n",
      "Train Epoch: 15 [53760/60000 (89%)]\tLoss: -0.751953\n",
      "Train Epoch: 15 [54272/60000 (90%)]\tLoss: -0.752487\n",
      "Train Epoch: 15 [54784/60000 (91%)]\tLoss: -0.757800\n",
      "Train Epoch: 15 [55296/60000 (92%)]\tLoss: -0.755194\n",
      "Train Epoch: 15 [55808/60000 (92%)]\tLoss: -0.756191\n",
      "Train Epoch: 15 [56320/60000 (93%)]\tLoss: -0.756909\n",
      "Train Epoch: 15 [56832/60000 (94%)]\tLoss: -0.755243\n",
      "Train Epoch: 15 [57344/60000 (95%)]\tLoss: -0.744684\n",
      "Train Epoch: 15 [57856/60000 (96%)]\tLoss: -0.745854\n",
      "Train Epoch: 15 [58368/60000 (97%)]\tLoss: -0.782185\n",
      "Train Epoch: 15 [58880/60000 (97%)]\tLoss: -0.757287\n",
      "Train Epoch: 15 [59392/60000 (98%)]\tLoss: -0.759984\n",
      "Train Epoch: 15 [11232/60000 (99%)]\tLoss: -3.767077\n",
      "====> Epoch: 15 Average loss: -0.7499\n",
      "====> Test set loss: -0.7850\n",
      "Train Epoch: 16 [0/60000 (0%)]\tLoss: -0.758777\n",
      "Train Epoch: 16 [512/60000 (1%)]\tLoss: -0.760692\n",
      "Train Epoch: 16 [1024/60000 (2%)]\tLoss: -0.747448\n",
      "Train Epoch: 16 [1536/60000 (3%)]\tLoss: -0.752189\n",
      "Train Epoch: 16 [2048/60000 (3%)]\tLoss: -0.795341\n",
      "Train Epoch: 16 [2560/60000 (4%)]\tLoss: -0.776570\n",
      "Train Epoch: 16 [3072/60000 (5%)]\tLoss: -0.737475\n",
      "Train Epoch: 16 [3584/60000 (6%)]\tLoss: -0.746600\n",
      "Train Epoch: 16 [4096/60000 (7%)]\tLoss: -0.754006\n",
      "Train Epoch: 16 [4608/60000 (8%)]\tLoss: -0.750250\n",
      "Train Epoch: 16 [5120/60000 (8%)]\tLoss: -0.772976\n",
      "Train Epoch: 16 [5632/60000 (9%)]\tLoss: -0.762797\n",
      "Train Epoch: 16 [6144/60000 (10%)]\tLoss: -0.736849\n",
      "Train Epoch: 16 [6656/60000 (11%)]\tLoss: -0.753623\n",
      "Train Epoch: 16 [7168/60000 (12%)]\tLoss: -0.745577\n",
      "Train Epoch: 16 [7680/60000 (13%)]\tLoss: -0.743050\n",
      "Train Epoch: 16 [8192/60000 (14%)]\tLoss: -0.789227\n",
      "Train Epoch: 16 [8704/60000 (14%)]\tLoss: -0.732080\n",
      "Train Epoch: 16 [9216/60000 (15%)]\tLoss: -0.748579\n",
      "Train Epoch: 16 [9728/60000 (16%)]\tLoss: -0.768814\n",
      "Train Epoch: 16 [10240/60000 (17%)]\tLoss: -0.780741\n",
      "Train Epoch: 16 [10752/60000 (18%)]\tLoss: -0.758658\n",
      "Train Epoch: 16 [11264/60000 (19%)]\tLoss: -0.784626\n",
      "Train Epoch: 16 [11776/60000 (19%)]\tLoss: -0.779125\n",
      "Train Epoch: 16 [12288/60000 (20%)]\tLoss: -0.770826\n",
      "Train Epoch: 16 [12800/60000 (21%)]\tLoss: -0.747148\n",
      "Train Epoch: 16 [13312/60000 (22%)]\tLoss: -0.772137\n",
      "Train Epoch: 16 [13824/60000 (23%)]\tLoss: -0.758057\n",
      "Train Epoch: 16 [14336/60000 (24%)]\tLoss: -0.756938\n",
      "Train Epoch: 16 [14848/60000 (25%)]\tLoss: -0.773846\n",
      "Train Epoch: 16 [15360/60000 (25%)]\tLoss: -0.681796\n",
      "Train Epoch: 16 [15872/60000 (26%)]\tLoss: -0.773806\n",
      "Train Epoch: 16 [16384/60000 (27%)]\tLoss: -0.767144\n",
      "Train Epoch: 16 [16896/60000 (28%)]\tLoss: -0.781279\n",
      "Train Epoch: 16 [17408/60000 (29%)]\tLoss: -0.768106\n",
      "Train Epoch: 16 [17920/60000 (30%)]\tLoss: -0.757604\n",
      "Train Epoch: 16 [18432/60000 (31%)]\tLoss: -0.760465\n",
      "Train Epoch: 16 [18944/60000 (31%)]\tLoss: -0.730128\n",
      "Train Epoch: 16 [19456/60000 (32%)]\tLoss: -0.750488\n",
      "Train Epoch: 16 [19968/60000 (33%)]\tLoss: -0.758319\n",
      "Train Epoch: 16 [20480/60000 (34%)]\tLoss: -0.748286\n",
      "Train Epoch: 16 [20992/60000 (35%)]\tLoss: -0.732467\n",
      "Train Epoch: 16 [21504/60000 (36%)]\tLoss: -0.754880\n",
      "Train Epoch: 16 [22016/60000 (36%)]\tLoss: -0.769227\n",
      "Train Epoch: 16 [22528/60000 (37%)]\tLoss: -0.760833\n",
      "Train Epoch: 16 [23040/60000 (38%)]\tLoss: -0.754128\n",
      "Train Epoch: 16 [23552/60000 (39%)]\tLoss: -0.766289\n",
      "Train Epoch: 16 [24064/60000 (40%)]\tLoss: -0.773029\n",
      "Train Epoch: 16 [24576/60000 (41%)]\tLoss: -0.792011\n",
      "Train Epoch: 16 [25088/60000 (42%)]\tLoss: -0.804854\n",
      "Train Epoch: 16 [25600/60000 (42%)]\tLoss: -0.767417\n",
      "Train Epoch: 16 [26112/60000 (43%)]\tLoss: -0.774087\n",
      "Train Epoch: 16 [26624/60000 (44%)]\tLoss: -0.764274\n",
      "Train Epoch: 16 [27136/60000 (45%)]\tLoss: -0.796427\n",
      "Train Epoch: 16 [27648/60000 (46%)]\tLoss: -0.759430\n",
      "Train Epoch: 16 [28160/60000 (47%)]\tLoss: -0.792057\n",
      "Train Epoch: 16 [28672/60000 (47%)]\tLoss: -0.775778\n",
      "Train Epoch: 16 [29184/60000 (48%)]\tLoss: -0.741321\n",
      "Train Epoch: 16 [29696/60000 (49%)]\tLoss: -0.759032\n",
      "Train Epoch: 16 [30208/60000 (50%)]\tLoss: -0.752444\n",
      "Train Epoch: 16 [30720/60000 (51%)]\tLoss: -0.763811\n",
      "Train Epoch: 16 [31232/60000 (52%)]\tLoss: -0.771207\n",
      "Train Epoch: 16 [31744/60000 (53%)]\tLoss: -0.753772\n",
      "Train Epoch: 16 [32256/60000 (53%)]\tLoss: -0.749989\n",
      "Train Epoch: 16 [32768/60000 (54%)]\tLoss: -0.751459\n",
      "Train Epoch: 16 [33280/60000 (55%)]\tLoss: -0.739951\n",
      "Train Epoch: 16 [33792/60000 (56%)]\tLoss: -0.768538\n",
      "Train Epoch: 16 [34304/60000 (57%)]\tLoss: -0.790332\n",
      "Train Epoch: 16 [34816/60000 (58%)]\tLoss: -0.723497\n",
      "Train Epoch: 16 [35328/60000 (58%)]\tLoss: -0.766213\n",
      "Train Epoch: 16 [35840/60000 (59%)]\tLoss: -0.771144\n",
      "Train Epoch: 16 [36352/60000 (60%)]\tLoss: -0.782469\n",
      "Train Epoch: 16 [36864/60000 (61%)]\tLoss: -0.798186\n",
      "Train Epoch: 16 [37376/60000 (62%)]\tLoss: -0.767692\n",
      "Train Epoch: 16 [37888/60000 (63%)]\tLoss: -0.756760\n",
      "Train Epoch: 16 [38400/60000 (64%)]\tLoss: -0.763896\n",
      "Train Epoch: 16 [38912/60000 (64%)]\tLoss: -0.787059\n",
      "Train Epoch: 16 [39424/60000 (65%)]\tLoss: -0.794394\n",
      "Train Epoch: 16 [39936/60000 (66%)]\tLoss: -0.768774\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train Epoch: 16 [40448/60000 (67%)]\tLoss: -0.773893\n",
      "Train Epoch: 16 [40960/60000 (68%)]\tLoss: -0.771169\n",
      "Train Epoch: 16 [41472/60000 (69%)]\tLoss: -0.761721\n",
      "Train Epoch: 16 [41984/60000 (69%)]\tLoss: -0.790034\n",
      "Train Epoch: 16 [42496/60000 (70%)]\tLoss: -0.763317\n",
      "Train Epoch: 16 [43008/60000 (71%)]\tLoss: -0.787948\n",
      "Train Epoch: 16 [43520/60000 (72%)]\tLoss: -0.783453\n",
      "Train Epoch: 16 [44032/60000 (73%)]\tLoss: -0.765876\n",
      "Train Epoch: 16 [44544/60000 (74%)]\tLoss: -0.776782\n",
      "Train Epoch: 16 [45056/60000 (75%)]\tLoss: -0.765851\n",
      "Train Epoch: 16 [45568/60000 (75%)]\tLoss: -0.731538\n",
      "Train Epoch: 16 [46080/60000 (76%)]\tLoss: -0.729736\n",
      "Train Epoch: 16 [46592/60000 (77%)]\tLoss: -0.776219\n",
      "Train Epoch: 16 [47104/60000 (78%)]\tLoss: -0.756522\n",
      "Train Epoch: 16 [47616/60000 (79%)]\tLoss: -0.764337\n",
      "Train Epoch: 16 [48128/60000 (80%)]\tLoss: -0.764561\n",
      "Train Epoch: 16 [48640/60000 (81%)]\tLoss: -0.756731\n",
      "Train Epoch: 16 [49152/60000 (81%)]\tLoss: -0.748479\n",
      "Train Epoch: 16 [49664/60000 (82%)]\tLoss: -0.786177\n",
      "Train Epoch: 16 [50176/60000 (83%)]\tLoss: -0.783463\n",
      "Train Epoch: 16 [50688/60000 (84%)]\tLoss: -0.764516\n",
      "Train Epoch: 16 [51200/60000 (85%)]\tLoss: -0.791284\n",
      "Train Epoch: 16 [51712/60000 (86%)]\tLoss: -0.706519\n",
      "Train Epoch: 16 [52224/60000 (86%)]\tLoss: -0.782604\n",
      "Train Epoch: 16 [52736/60000 (87%)]\tLoss: -0.749759\n",
      "Train Epoch: 16 [53248/60000 (88%)]\tLoss: -0.779235\n",
      "Train Epoch: 16 [53760/60000 (89%)]\tLoss: -0.770303\n",
      "Train Epoch: 16 [54272/60000 (90%)]\tLoss: -0.753868\n",
      "Train Epoch: 16 [54784/60000 (91%)]\tLoss: -0.774813\n",
      "Train Epoch: 16 [55296/60000 (92%)]\tLoss: -0.750973\n",
      "Train Epoch: 16 [55808/60000 (92%)]\tLoss: -0.755084\n",
      "Train Epoch: 16 [56320/60000 (93%)]\tLoss: -0.602649\n",
      "Train Epoch: 16 [56832/60000 (94%)]\tLoss: -0.760242\n",
      "Train Epoch: 16 [57344/60000 (95%)]\tLoss: -0.763835\n",
      "Train Epoch: 16 [57856/60000 (96%)]\tLoss: -0.743077\n",
      "Train Epoch: 16 [58368/60000 (97%)]\tLoss: -0.752909\n",
      "Train Epoch: 16 [58880/60000 (97%)]\tLoss: -0.750642\n",
      "Train Epoch: 16 [59392/60000 (98%)]\tLoss: -0.744608\n",
      "Train Epoch: 16 [11232/60000 (99%)]\tLoss: -4.358452\n",
      "====> Epoch: 16 Average loss: -0.7672\n",
      "====> Test set loss: -0.7907\n",
      "Train Epoch: 17 [0/60000 (0%)]\tLoss: -0.744235\n",
      "Train Epoch: 17 [512/60000 (1%)]\tLoss: -0.761981\n",
      "Train Epoch: 17 [1024/60000 (2%)]\tLoss: -0.702317\n",
      "Train Epoch: 17 [1536/60000 (3%)]\tLoss: -0.755170\n",
      "Train Epoch: 17 [2048/60000 (3%)]\tLoss: -0.774890\n",
      "Train Epoch: 17 [2560/60000 (4%)]\tLoss: -0.664174\n",
      "Train Epoch: 17 [3072/60000 (5%)]\tLoss: -0.749476\n",
      "Train Epoch: 17 [3584/60000 (6%)]\tLoss: -0.769208\n",
      "Train Epoch: 17 [4096/60000 (7%)]\tLoss: -0.770176\n",
      "Train Epoch: 17 [4608/60000 (8%)]\tLoss: -0.756843\n",
      "Train Epoch: 17 [5120/60000 (8%)]\tLoss: -0.781461\n",
      "Train Epoch: 17 [5632/60000 (9%)]\tLoss: -0.523379\n",
      "Train Epoch: 17 [6144/60000 (10%)]\tLoss: -0.789768\n",
      "Train Epoch: 17 [6656/60000 (11%)]\tLoss: -0.772727\n",
      "Train Epoch: 17 [7168/60000 (12%)]\tLoss: -0.752670\n",
      "Train Epoch: 17 [7680/60000 (13%)]\tLoss: -0.667683\n",
      "Train Epoch: 17 [8192/60000 (14%)]\tLoss: -0.740165\n",
      "Train Epoch: 17 [8704/60000 (14%)]\tLoss: -0.715267\n",
      "Train Epoch: 17 [9216/60000 (15%)]\tLoss: -0.435478\n",
      "Train Epoch: 17 [9728/60000 (16%)]\tLoss: -0.754039\n",
      "Train Epoch: 17 [10240/60000 (17%)]\tLoss: -0.776435\n",
      "Train Epoch: 17 [10752/60000 (18%)]\tLoss: -0.745519\n",
      "Train Epoch: 17 [11264/60000 (19%)]\tLoss: -0.639888\n",
      "Train Epoch: 17 [11776/60000 (19%)]\tLoss: -0.740076\n",
      "Train Epoch: 17 [12288/60000 (20%)]\tLoss: -0.724947\n",
      "Train Epoch: 17 [12800/60000 (21%)]\tLoss: -0.746203\n",
      "Train Epoch: 17 [13312/60000 (22%)]\tLoss: -0.737698\n",
      "Train Epoch: 17 [13824/60000 (23%)]\tLoss: -0.749255\n",
      "Train Epoch: 17 [14336/60000 (24%)]\tLoss: -0.766786\n",
      "Train Epoch: 17 [14848/60000 (25%)]\tLoss: -0.749794\n",
      "Train Epoch: 17 [15360/60000 (25%)]\tLoss: -0.757926\n",
      "Train Epoch: 17 [15872/60000 (26%)]\tLoss: -0.779439\n",
      "Train Epoch: 17 [16384/60000 (27%)]\tLoss: -0.795411\n",
      "Train Epoch: 17 [16896/60000 (28%)]\tLoss: -0.780900\n",
      "Train Epoch: 17 [17408/60000 (29%)]\tLoss: -0.772049\n",
      "Train Epoch: 17 [17920/60000 (30%)]\tLoss: -0.763580\n",
      "Train Epoch: 17 [18432/60000 (31%)]\tLoss: -0.752649\n",
      "Train Epoch: 17 [18944/60000 (31%)]\tLoss: -0.749878\n",
      "Train Epoch: 17 [19456/60000 (32%)]\tLoss: -0.766467\n",
      "Train Epoch: 17 [19968/60000 (33%)]\tLoss: -0.769515\n",
      "Train Epoch: 17 [20480/60000 (34%)]\tLoss: -0.758833\n",
      "Train Epoch: 17 [20992/60000 (35%)]\tLoss: -0.782249\n",
      "Train Epoch: 17 [21504/60000 (36%)]\tLoss: -0.766278\n",
      "Train Epoch: 17 [22016/60000 (36%)]\tLoss: -0.777203\n",
      "Train Epoch: 17 [22528/60000 (37%)]\tLoss: -0.734008\n",
      "Train Epoch: 17 [23040/60000 (38%)]\tLoss: -0.772170\n",
      "Train Epoch: 17 [23552/60000 (39%)]\tLoss: -0.778563\n",
      "Train Epoch: 17 [24064/60000 (40%)]\tLoss: -0.752750\n",
      "Train Epoch: 17 [24576/60000 (41%)]\tLoss: -0.789940\n",
      "Train Epoch: 17 [25088/60000 (42%)]\tLoss: -0.674965\n",
      "Train Epoch: 17 [25600/60000 (42%)]\tLoss: -0.742251\n",
      "Train Epoch: 17 [26112/60000 (43%)]\tLoss: -0.764286\n",
      "Train Epoch: 17 [26624/60000 (44%)]\tLoss: -0.745767\n",
      "Train Epoch: 17 [27136/60000 (45%)]\tLoss: -0.746999\n",
      "Train Epoch: 17 [27648/60000 (46%)]\tLoss: -0.766879\n",
      "Train Epoch: 17 [28160/60000 (47%)]\tLoss: -0.784550\n",
      "Train Epoch: 17 [28672/60000 (47%)]\tLoss: -0.765307\n",
      "Train Epoch: 17 [29184/60000 (48%)]\tLoss: -0.757485\n",
      "Train Epoch: 17 [29696/60000 (49%)]\tLoss: -0.705171\n",
      "Train Epoch: 17 [30208/60000 (50%)]\tLoss: -0.778992\n",
      "Train Epoch: 17 [30720/60000 (51%)]\tLoss: -0.742528\n",
      "Train Epoch: 17 [31232/60000 (52%)]\tLoss: -0.769966\n",
      "Train Epoch: 17 [31744/60000 (53%)]\tLoss: -0.712693\n",
      "Train Epoch: 17 [32256/60000 (53%)]\tLoss: -0.767823\n",
      "Train Epoch: 17 [32768/60000 (54%)]\tLoss: -0.798584\n",
      "Train Epoch: 17 [33280/60000 (55%)]\tLoss: -0.746750\n",
      "Train Epoch: 17 [33792/60000 (56%)]\tLoss: -0.753652\n",
      "Train Epoch: 17 [34304/60000 (57%)]\tLoss: -0.773732\n",
      "Train Epoch: 17 [34816/60000 (58%)]\tLoss: -0.786362\n",
      "Train Epoch: 17 [35328/60000 (58%)]\tLoss: -0.763691\n",
      "Train Epoch: 17 [35840/60000 (59%)]\tLoss: -0.765856\n",
      "Train Epoch: 17 [36352/60000 (60%)]\tLoss: -0.728097\n",
      "Train Epoch: 17 [36864/60000 (61%)]\tLoss: -0.756493\n",
      "Train Epoch: 17 [37376/60000 (62%)]\tLoss: -0.592285\n",
      "Train Epoch: 17 [37888/60000 (63%)]\tLoss: -0.688532\n",
      "Train Epoch: 17 [38400/60000 (64%)]\tLoss: -0.757056\n",
      "Train Epoch: 17 [38912/60000 (64%)]\tLoss: -0.746599\n",
      "Train Epoch: 17 [39424/60000 (65%)]\tLoss: -0.754805\n",
      "Train Epoch: 17 [39936/60000 (66%)]\tLoss: -0.731060\n",
      "Train Epoch: 17 [40448/60000 (67%)]\tLoss: -0.761563\n",
      "Train Epoch: 17 [40960/60000 (68%)]\tLoss: -0.748460\n",
      "Train Epoch: 17 [41472/60000 (69%)]\tLoss: -0.719282\n",
      "Train Epoch: 17 [41984/60000 (69%)]\tLoss: -0.764808\n",
      "Train Epoch: 17 [42496/60000 (70%)]\tLoss: -0.748404\n",
      "Train Epoch: 17 [43008/60000 (71%)]\tLoss: -0.767694\n",
      "Train Epoch: 17 [43520/60000 (72%)]\tLoss: -0.749802\n",
      "Train Epoch: 17 [44032/60000 (73%)]\tLoss: -0.716070\n",
      "Train Epoch: 17 [44544/60000 (74%)]\tLoss: -0.778523\n",
      "Train Epoch: 17 [45056/60000 (75%)]\tLoss: -0.739315\n",
      "Train Epoch: 17 [45568/60000 (75%)]\tLoss: -0.759551\n",
      "Train Epoch: 17 [46080/60000 (76%)]\tLoss: -0.766719\n",
      "Train Epoch: 17 [46592/60000 (77%)]\tLoss: -0.746700\n",
      "Train Epoch: 17 [47104/60000 (78%)]\tLoss: -0.766068\n",
      "Train Epoch: 17 [47616/60000 (79%)]\tLoss: -0.750319\n",
      "Train Epoch: 17 [48128/60000 (80%)]\tLoss: -0.759081\n",
      "Train Epoch: 17 [48640/60000 (81%)]\tLoss: -0.744854\n",
      "Train Epoch: 17 [49152/60000 (81%)]\tLoss: -0.747079\n",
      "Train Epoch: 17 [49664/60000 (82%)]\tLoss: -0.808727\n",
      "Train Epoch: 17 [50176/60000 (83%)]\tLoss: -0.770594\n",
      "Train Epoch: 17 [50688/60000 (84%)]\tLoss: -0.664843\n",
      "Train Epoch: 17 [51200/60000 (85%)]\tLoss: -0.763610\n",
      "Train Epoch: 17 [51712/60000 (86%)]\tLoss: -0.766589\n",
      "Train Epoch: 17 [52224/60000 (86%)]\tLoss: -0.756479\n",
      "Train Epoch: 17 [52736/60000 (87%)]\tLoss: -0.740528\n",
      "Train Epoch: 17 [53248/60000 (88%)]\tLoss: -0.739583\n",
      "Train Epoch: 17 [53760/60000 (89%)]\tLoss: -0.753406\n",
      "Train Epoch: 17 [54272/60000 (90%)]\tLoss: -0.747611\n",
      "Train Epoch: 17 [54784/60000 (91%)]\tLoss: -0.778981\n",
      "Train Epoch: 17 [55296/60000 (92%)]\tLoss: -0.728957\n",
      "Train Epoch: 17 [55808/60000 (92%)]\tLoss: -0.749095\n",
      "Train Epoch: 17 [56320/60000 (93%)]\tLoss: -0.763731\n",
      "Train Epoch: 17 [56832/60000 (94%)]\tLoss: -0.721404\n",
      "Train Epoch: 17 [57344/60000 (95%)]\tLoss: -0.748220\n",
      "Train Epoch: 17 [57856/60000 (96%)]\tLoss: -0.757951\n",
      "Train Epoch: 17 [58368/60000 (97%)]\tLoss: -0.784842\n",
      "Train Epoch: 17 [58880/60000 (97%)]\tLoss: -0.758681\n",
      "Train Epoch: 17 [59392/60000 (98%)]\tLoss: -0.767017\n",
      "Train Epoch: 17 [11232/60000 (99%)]\tLoss: -4.227625\n",
      "====> Epoch: 17 Average loss: -0.7520\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "====> Test set loss: -0.7843\n",
      "Train Epoch: 18 [0/60000 (0%)]\tLoss: -0.752082\n",
      "Train Epoch: 18 [512/60000 (1%)]\tLoss: -0.729629\n",
      "Train Epoch: 18 [1024/60000 (2%)]\tLoss: -0.733355\n",
      "Train Epoch: 18 [1536/60000 (3%)]\tLoss: -0.751919\n",
      "Train Epoch: 18 [2048/60000 (3%)]\tLoss: -0.760740\n",
      "Train Epoch: 18 [2560/60000 (4%)]\tLoss: -0.753886\n",
      "Train Epoch: 18 [3072/60000 (5%)]\tLoss: -0.762269\n",
      "Train Epoch: 18 [3584/60000 (6%)]\tLoss: -0.710559\n",
      "Train Epoch: 18 [4096/60000 (7%)]\tLoss: -0.777625\n",
      "Train Epoch: 18 [4608/60000 (8%)]\tLoss: -0.728755\n",
      "Train Epoch: 18 [5120/60000 (8%)]\tLoss: -0.761260\n",
      "Train Epoch: 18 [5632/60000 (9%)]\tLoss: -0.680493\n",
      "Train Epoch: 18 [6144/60000 (10%)]\tLoss: -0.749966\n",
      "Train Epoch: 18 [6656/60000 (11%)]\tLoss: -0.791898\n",
      "Train Epoch: 18 [7168/60000 (12%)]\tLoss: -0.786334\n",
      "Train Epoch: 18 [7680/60000 (13%)]\tLoss: -0.764163\n",
      "Train Epoch: 18 [8192/60000 (14%)]\tLoss: -0.782414\n",
      "Train Epoch: 18 [8704/60000 (14%)]\tLoss: -0.756358\n",
      "Train Epoch: 18 [9216/60000 (15%)]\tLoss: -0.755179\n",
      "Train Epoch: 18 [9728/60000 (16%)]\tLoss: -0.528698\n",
      "Train Epoch: 18 [10240/60000 (17%)]\tLoss: -0.702290\n",
      "Train Epoch: 18 [10752/60000 (18%)]\tLoss: -0.795817\n",
      "Train Epoch: 18 [11264/60000 (19%)]\tLoss: -0.792151\n",
      "Train Epoch: 18 [11776/60000 (19%)]\tLoss: -0.778075\n",
      "Train Epoch: 18 [12288/60000 (20%)]\tLoss: -0.769994\n",
      "Train Epoch: 18 [12800/60000 (21%)]\tLoss: -0.701445\n",
      "Train Epoch: 18 [13312/60000 (22%)]\tLoss: -0.786349\n",
      "Train Epoch: 18 [13824/60000 (23%)]\tLoss: -0.765691\n",
      "Train Epoch: 18 [14336/60000 (24%)]\tLoss: -0.761428\n",
      "Train Epoch: 18 [14848/60000 (25%)]\tLoss: -0.781839\n",
      "Train Epoch: 18 [15360/60000 (25%)]\tLoss: -0.757389\n",
      "Train Epoch: 18 [15872/60000 (26%)]\tLoss: -0.758222\n",
      "Train Epoch: 18 [16384/60000 (27%)]\tLoss: -0.748501\n",
      "Train Epoch: 18 [16896/60000 (28%)]\tLoss: -0.766799\n",
      "Train Epoch: 18 [17408/60000 (29%)]\tLoss: -0.774531\n",
      "Train Epoch: 18 [17920/60000 (30%)]\tLoss: -0.790196\n",
      "Train Epoch: 18 [18432/60000 (31%)]\tLoss: -0.757317\n",
      "Train Epoch: 18 [18944/60000 (31%)]\tLoss: -0.794558\n",
      "Train Epoch: 18 [19456/60000 (32%)]\tLoss: -0.755399\n",
      "Train Epoch: 18 [19968/60000 (33%)]\tLoss: -0.762251\n",
      "Train Epoch: 18 [20480/60000 (34%)]\tLoss: -0.764550\n",
      "Train Epoch: 18 [20992/60000 (35%)]\tLoss: -0.703552\n",
      "Train Epoch: 18 [21504/60000 (36%)]\tLoss: -0.744123\n",
      "Train Epoch: 18 [22016/60000 (36%)]\tLoss: -0.764117\n",
      "Train Epoch: 18 [22528/60000 (37%)]\tLoss: -0.772502\n",
      "Train Epoch: 18 [23040/60000 (38%)]\tLoss: -0.752529\n",
      "Train Epoch: 18 [23552/60000 (39%)]\tLoss: -0.785031\n",
      "Train Epoch: 18 [24064/60000 (40%)]\tLoss: -0.690513\n",
      "Train Epoch: 18 [24576/60000 (41%)]\tLoss: -0.762408\n",
      "Train Epoch: 18 [25088/60000 (42%)]\tLoss: -0.739639\n",
      "Train Epoch: 18 [25600/60000 (42%)]\tLoss: -0.758912\n",
      "Train Epoch: 18 [26112/60000 (43%)]\tLoss: -0.583391\n",
      "Train Epoch: 18 [26624/60000 (44%)]\tLoss: -0.775218\n",
      "Train Epoch: 18 [27136/60000 (45%)]\tLoss: -0.766913\n",
      "Train Epoch: 18 [27648/60000 (46%)]\tLoss: -0.735006\n",
      "Train Epoch: 18 [28160/60000 (47%)]\tLoss: -0.761111\n",
      "Train Epoch: 18 [28672/60000 (47%)]\tLoss: -0.778999\n",
      "Train Epoch: 18 [29184/60000 (48%)]\tLoss: -0.751682\n",
      "Train Epoch: 18 [29696/60000 (49%)]\tLoss: -0.756065\n",
      "Train Epoch: 18 [30208/60000 (50%)]\tLoss: -0.750097\n",
      "Train Epoch: 18 [30720/60000 (51%)]\tLoss: -0.753913\n",
      "Train Epoch: 18 [31232/60000 (52%)]\tLoss: -0.771373\n",
      "Train Epoch: 18 [31744/60000 (53%)]\tLoss: -0.778640\n",
      "Train Epoch: 18 [32256/60000 (53%)]\tLoss: -0.737850\n",
      "Train Epoch: 18 [32768/60000 (54%)]\tLoss: -0.770928\n",
      "Train Epoch: 18 [33280/60000 (55%)]\tLoss: -0.739402\n",
      "Train Epoch: 18 [33792/60000 (56%)]\tLoss: -0.802108\n",
      "Train Epoch: 18 [34304/60000 (57%)]\tLoss: -0.766064\n",
      "Train Epoch: 18 [34816/60000 (58%)]\tLoss: -0.763519\n",
      "Train Epoch: 18 [35328/60000 (58%)]\tLoss: -0.777173\n",
      "Train Epoch: 18 [35840/60000 (59%)]\tLoss: -0.785232\n",
      "Train Epoch: 18 [36352/60000 (60%)]\tLoss: -0.776630\n",
      "Train Epoch: 18 [36864/60000 (61%)]\tLoss: -0.760570\n",
      "Train Epoch: 18 [37376/60000 (62%)]\tLoss: -0.763204\n",
      "Train Epoch: 18 [37888/60000 (63%)]\tLoss: -0.742453\n",
      "Train Epoch: 18 [38400/60000 (64%)]\tLoss: -0.763922\n",
      "Train Epoch: 18 [38912/60000 (64%)]\tLoss: -0.714577\n",
      "Train Epoch: 18 [39424/60000 (65%)]\tLoss: -0.770157\n",
      "Train Epoch: 18 [39936/60000 (66%)]\tLoss: -0.763524\n",
      "Train Epoch: 18 [40448/60000 (67%)]\tLoss: -0.776317\n",
      "Train Epoch: 18 [40960/60000 (68%)]\tLoss: -0.744119\n",
      "Train Epoch: 18 [41472/60000 (69%)]\tLoss: -0.672805\n",
      "Train Epoch: 18 [41984/60000 (69%)]\tLoss: -0.752190\n",
      "Train Epoch: 18 [42496/60000 (70%)]\tLoss: -0.759733\n",
      "Train Epoch: 18 [43008/60000 (71%)]\tLoss: -0.776050\n",
      "Train Epoch: 18 [43520/60000 (72%)]\tLoss: -0.765773\n",
      "Train Epoch: 18 [44032/60000 (73%)]\tLoss: -0.748575\n",
      "Train Epoch: 18 [44544/60000 (74%)]\tLoss: -0.726078\n",
      "Train Epoch: 18 [45056/60000 (75%)]\tLoss: -0.764468\n",
      "Train Epoch: 18 [45568/60000 (75%)]\tLoss: -0.724342\n",
      "Train Epoch: 18 [46080/60000 (76%)]\tLoss: -0.729774\n",
      "Train Epoch: 18 [46592/60000 (77%)]\tLoss: -0.657675\n",
      "Train Epoch: 18 [47104/60000 (78%)]\tLoss: -0.736312\n",
      "Train Epoch: 18 [47616/60000 (79%)]\tLoss: -0.747089\n",
      "Train Epoch: 18 [48128/60000 (80%)]\tLoss: -0.741098\n",
      "Train Epoch: 18 [48640/60000 (81%)]\tLoss: -0.728666\n",
      "Train Epoch: 18 [49152/60000 (81%)]\tLoss: -0.778705\n",
      "Train Epoch: 18 [49664/60000 (82%)]\tLoss: -0.764329\n",
      "Train Epoch: 18 [50176/60000 (83%)]\tLoss: -0.710884\n",
      "Train Epoch: 18 [50688/60000 (84%)]\tLoss: -0.738491\n",
      "Train Epoch: 18 [51200/60000 (85%)]\tLoss: -0.730378\n",
      "Train Epoch: 18 [51712/60000 (86%)]\tLoss: -0.752663\n",
      "Train Epoch: 18 [52224/60000 (86%)]\tLoss: -0.767972\n",
      "Train Epoch: 18 [52736/60000 (87%)]\tLoss: -0.578274\n",
      "Train Epoch: 18 [53248/60000 (88%)]\tLoss: -0.804917\n",
      "Train Epoch: 18 [53760/60000 (89%)]\tLoss: -0.766978\n",
      "Train Epoch: 18 [54272/60000 (90%)]\tLoss: -0.769180\n",
      "Train Epoch: 18 [54784/60000 (91%)]\tLoss: -0.767465\n",
      "Train Epoch: 18 [55296/60000 (92%)]\tLoss: -0.761366\n",
      "Train Epoch: 18 [55808/60000 (92%)]\tLoss: -0.715730\n",
      "Train Epoch: 18 [56320/60000 (93%)]\tLoss: -0.770185\n",
      "Train Epoch: 18 [56832/60000 (94%)]\tLoss: -0.788979\n",
      "Train Epoch: 18 [57344/60000 (95%)]\tLoss: -0.747381\n",
      "Train Epoch: 18 [57856/60000 (96%)]\tLoss: -0.725300\n",
      "Train Epoch: 18 [58368/60000 (97%)]\tLoss: -0.742365\n",
      "Train Epoch: 18 [58880/60000 (97%)]\tLoss: -0.732517\n",
      "Train Epoch: 18 [59392/60000 (98%)]\tLoss: -0.721862\n",
      "Train Epoch: 18 [11232/60000 (99%)]\tLoss: -3.924689\n",
      "====> Epoch: 18 Average loss: -0.7547\n",
      "====> Test set loss: -0.7868\n",
      "Train Epoch: 19 [0/60000 (0%)]\tLoss: -0.764440\n",
      "Train Epoch: 19 [512/60000 (1%)]\tLoss: -0.774185\n",
      "Train Epoch: 19 [1024/60000 (2%)]\tLoss: -0.713171\n",
      "Train Epoch: 19 [1536/60000 (3%)]\tLoss: -0.774133\n",
      "Train Epoch: 19 [2048/60000 (3%)]\tLoss: -0.753963\n",
      "Train Epoch: 19 [2560/60000 (4%)]\tLoss: -0.771722\n",
      "Train Epoch: 19 [3072/60000 (5%)]\tLoss: -0.778077\n",
      "Train Epoch: 19 [3584/60000 (6%)]\tLoss: -0.751113\n",
      "Train Epoch: 19 [4096/60000 (7%)]\tLoss: -0.754684\n",
      "Train Epoch: 19 [4608/60000 (8%)]\tLoss: -0.757954\n",
      "Train Epoch: 19 [5120/60000 (8%)]\tLoss: -0.747602\n",
      "Train Epoch: 19 [5632/60000 (9%)]\tLoss: -0.753818\n",
      "Train Epoch: 19 [6144/60000 (10%)]\tLoss: -0.787809\n",
      "Train Epoch: 19 [6656/60000 (11%)]\tLoss: -0.783520\n",
      "Train Epoch: 19 [7168/60000 (12%)]\tLoss: -0.740272\n",
      "Train Epoch: 19 [7680/60000 (13%)]\tLoss: -0.735787\n",
      "Train Epoch: 19 [8192/60000 (14%)]\tLoss: -0.742363\n",
      "Train Epoch: 19 [8704/60000 (14%)]\tLoss: -0.747140\n",
      "Train Epoch: 19 [9216/60000 (15%)]\tLoss: -0.772598\n",
      "Train Epoch: 19 [9728/60000 (16%)]\tLoss: -0.752984\n",
      "Train Epoch: 19 [10240/60000 (17%)]\tLoss: -0.778512\n",
      "Train Epoch: 19 [10752/60000 (18%)]\tLoss: -0.765337\n",
      "Train Epoch: 19 [11264/60000 (19%)]\tLoss: -0.734108\n",
      "Train Epoch: 19 [11776/60000 (19%)]\tLoss: -0.711208\n",
      "Train Epoch: 19 [12288/60000 (20%)]\tLoss: -0.738664\n",
      "Train Epoch: 19 [12800/60000 (21%)]\tLoss: -0.760782\n"
     ]
    }
   ],
   "source": [
    "for epoch in range(14, epochs + 1):\n",
    "    train(epoch)\n",
    "    test(epoch)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.save(model.state_dict(), result_path+\"model.pt\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "optimizer = optim.Adam(model.parameters(), lr=1e-4)\n",
    "epochs=500"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "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.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}