Capsule Network Train.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import numpy\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from tqdm import tqdm"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Model Loading"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "USE_CUDA = True if torch.cuda.is_available() else False\n",
    "\n",
    "class ConvLayer(nn.Module):\n",
    "    def __init__(self, in_channels=1, out_channels=256, kernel_size=9):\n",
    "        super(ConvLayer, self).__init__()\n",
    "\n",
    "        self.conv = nn.Conv2d(in_channels=in_channels,\n",
    "                              out_channels=out_channels,\n",
    "                              kernel_size=kernel_size,\n",
    "                              stride=1\n",
    "                              )\n",
    "\n",
    "    def forward(self, x):\n",
    "        return F.relu(self.conv(x))\n",
    "\n",
    "\n",
    "class PrimaryCaps(nn.Module):\n",
    "    def __init__(self, num_capsules=8, in_channels=256, out_channels=32, kernel_size=9, num_routes=32 * 6 * 6):\n",
    "        super(PrimaryCaps, self).__init__()\n",
    "        self.num_routes = num_routes\n",
    "        self.capsules = nn.ModuleList([\n",
    "            nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=kernel_size, stride=1, padding=0)\n",
    "            for _ in range(num_capsules)])\n",
    "\n",
    "    def forward(self, x):\n",
    "        u = [capsule(x) for capsule in self.capsules]        \n",
    "        u = torch.stack(u, dim=1)\n",
    "        u = u.view(x.size(0), self.num_routes, -1)\n",
    "        return self.squash(u)\n",
    "\n",
    "    def squash(self, input_tensor, epsilon=1e-7):\n",
    "        squared_norm = (input_tensor ** 2 + epsilon).sum(-1, keepdim=True)\n",
    "        output_tensor = squared_norm * input_tensor / ((1. + squared_norm) * torch.sqrt(squared_norm))\n",
    "        return output_tensor\n",
    "\n",
    "\n",
    "class DigitCaps(nn.Module):\n",
    "    def __init__(self, num_capsules=10, num_routes=32 * 6 * 6, in_channels=8, out_channels=16):\n",
    "        super(DigitCaps, self).__init__()\n",
    "\n",
    "        self.in_channels = in_channels\n",
    "        self.num_routes = num_routes\n",
    "        self.num_capsules = num_capsules\n",
    "\n",
    "        self.W = nn.Parameter(torch.randn(1, num_routes, num_capsules, out_channels, in_channels))\n",
    "\n",
    "    def forward(self, x):\n",
    "        batch_size = x.size(0)\n",
    "        x = torch.stack([x] * self.num_capsules, dim=2).unsqueeze(4)\n",
    "\n",
    "        W = torch.cat([self.W] * batch_size, dim=0)\n",
    "        u_hat = torch.matmul(W, x)\n",
    "\n",
    "        b_ij = torch.zeros(1, self.num_routes, self.num_capsules, 1)\n",
    "        if USE_CUDA:\n",
    "            b_ij = b_ij.cuda()\n",
    "\n",
    "        num_iterations = 3\n",
    "        for iteration in range(num_iterations):\n",
    "            c_ij = F.softmax(b_ij, dim=1)\n",
    "            c_ij = torch.cat([c_ij] * batch_size, dim=0).unsqueeze(4)\n",
    "\n",
    "            s_j = (c_ij * u_hat).sum(dim=1, keepdim=True)\n",
    "            v_j = self.squash(s_j)\n",
    "\n",
    "            if iteration < num_iterations - 1:\n",
    "                a_ij = torch.matmul(u_hat.transpose(3, 4), torch.cat([v_j] * self.num_routes, dim=1))\n",
    "                b_ij = b_ij + a_ij.squeeze(4).mean(dim=0, keepdim=True)\n",
    "\n",
    "        return v_j.squeeze(1)\n",
    "\n",
    "    def squash(self, input_tensor, epsilon=1e-7):\n",
    "        squared_norm = (input_tensor ** 2 + epsilon).sum(-1, keepdim=True)\n",
    "        output_tensor = squared_norm * input_tensor / ((1. + squared_norm) * torch.sqrt(squared_norm))\n",
    "        return output_tensor\n",
    "\n",
    "\n",
    "class Decoder(nn.Module):\n",
    "    def __init__(self, input_width=28, input_height=28, input_channel=1):\n",
    "        super(Decoder, self).__init__()\n",
    "        self.input_width = input_width\n",
    "        self.input_height = input_height\n",
    "        self.input_channel = input_channel\n",
    "        self.reconstraction_layers = nn.Sequential(\n",
    "            nn.Linear(16 * 10, 512),\n",
    "            nn.ReLU(inplace=True),\n",
    "            nn.Linear(512, 1024),\n",
    "            nn.ReLU(inplace=True),\n",
    "            nn.Linear(1024, self.input_height * self.input_height * self.input_channel),\n",
    "            nn.Sigmoid()\n",
    "        )\n",
    "\n",
    "    def forward(self, x, data):\n",
    "        classes = torch.sqrt((x ** 2).sum(2))\n",
    "        classes = F.softmax(classes.squeeze(), dim=1)\n",
    "\n",
    "        _, max_length_indices = classes.max(dim=1)\n",
    "        masked = torch.sparse.torch.eye(10)\n",
    "        if USE_CUDA:\n",
    "            masked = masked.cuda()\n",
    "        masked = masked.index_select(dim=0, index=max_length_indices.squeeze().data)\n",
    "        t = (x * masked[:, :, None, None]).view(x.size(0), -1)\n",
    "        reconstructions = self.reconstraction_layers(t)\n",
    "        reconstructions = reconstructions.view(-1, self.input_channel, self.input_width, self.input_height)\n",
    "        return reconstructions, masked\n",
    "\n",
    "\n",
    "class CapsNet(nn.Module):\n",
    "    def __init__(self, config=None):\n",
    "        super(CapsNet, self).__init__()\n",
    "        if config:\n",
    "            self.conv_layer = ConvLayer(config.cnn_in_channels, config.cnn_out_channels, config.cnn_kernel_size)\n",
    "            self.primary_capsules = PrimaryCaps(config.pc_num_capsules, config.pc_in_channels, config.pc_out_channels,\n",
    "                                                config.pc_kernel_size, config.pc_num_routes)\n",
    "            self.digit_capsules_1 = DigitCaps(config.dc_num_capsules, config.dc_num_routes, config.dc_in_channels,\n",
    "                                            config.dc_out_channels)\n",
    "            self.digit_capsules_2 = DigitCaps(config.dc_2_num_capsules, config.dc_2_num_routes, config.dc_2_in_channels,\n",
    "                                            config.dc_2_out_channels)\n",
    "            self.decoder = Decoder(config.input_width, config.input_height, config.cnn_in_channels)\n",
    "        else:\n",
    "            self.conv_layer = ConvLayer()\n",
    "            self.primary_capsules = PrimaryCaps()\n",
    "            self.digit_capsules_1 = DigitCaps()\n",
    "            self.decoder = Decoder()\n",
    "\n",
    "        self.mse_loss = nn.MSELoss()\n",
    "\n",
    "    def forward(self, data):\n",
    "        output = self.digit_capsules_1(self.primary_capsules(self.conv_layer(data))).squeeze(-1)\n",
    "        output = self.digit_capsules_2(output)\n",
    "        reconstructions, masked = self.decoder(output, data)\n",
    "        return output, reconstructions, masked\n",
    "\n",
    "    def loss(self, data, x, target, reconstructions):\n",
    "        return self.margin_loss(x, target) + self.reconstruction_loss(data, reconstructions)\n",
    "\n",
    "    def margin_loss(self, x, labels, size_average=True):\n",
    "        batch_size = x.size(0)\n",
    "\n",
    "        v_c = torch.sqrt((x ** 2).sum(dim=2, keepdim=True))\n",
    "\n",
    "        left = F.relu(0.9 - v_c).view(batch_size, -1)\n",
    "        right = F.relu(v_c - 0.1).view(batch_size, -1)\n",
    "\n",
    "        loss = labels * left + 0.5 * (1.0 - labels) * right\n",
    "        loss = loss.sum(dim=1).mean()\n",
    "\n",
    "        return loss\n",
    "\n",
    "    def reconstruction_loss(self, data, reconstructions):\n",
    "        loss = self.mse_loss(reconstructions.view(reconstructions.size(0), -1), data.view(reconstructions.size(0), -1))\n",
    "        return loss * 0.0005"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Config for 49 16d vectors in the Primary Capsule. Set Softmax dimension to 0 in this case\n",
    "class Config:\n",
    "    def __init__(self, dataset='mnist'):\n",
    "        # CNN (cnn)\n",
    "        self.cnn_in_channels = 1\n",
    "        self.cnn_out_channels = 12\n",
    "        self.cnn_kernel_size = 15\n",
    "\n",
    "        # Primary Capsule (pc)\n",
    "        self.pc_num_capsules = 16\n",
    "        self.pc_in_channels = 12\n",
    "        self.pc_out_channels = 1\n",
    "        self.pc_kernel_size = 8\n",
    "        self.pc_num_routes = 1 * 7 * 7\n",
    "\n",
    "        # Digit Capsule 1 (dc)\n",
    "        self.dc_num_capsules = 49\n",
    "        self.dc_num_routes = 1 * 7 * 7\n",
    "        self.dc_in_channels = 16\n",
    "        self.dc_out_channels = 16\n",
    "        \n",
    "        # Digit Capsule 2 (dc)\n",
    "        self.dc_2_num_capsules = 10\n",
    "        self.dc_2_num_routes = 1 * 7 * 7\n",
    "        self.dc_2_in_channels = 16\n",
    "        self.dc_2_out_channels = 16\n",
    "\n",
    "        # Decoder\n",
    "        self.input_width = 28\n",
    "        self.input_height = 28"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Training CapsuleNet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "USE_CUDA = True if torch.cuda.is_available() else False\n",
    "BATCH_SIZE = 256\n",
    "N_EPOCHS = 30\n",
    "LEARNING_RATE = 0.01\n",
    "MOMENTUM = 0.9\n",
    "\n",
    "def train(model, optimizer, train_loader, epoch):\n",
    "    capsule_net = model\n",
    "    capsule_net.train()\n",
    "    n_batch = len(train_loader)\n",
    "    total_loss = 0\n",
    "    for batch_id, (data, target) in enumerate(tqdm(train_loader)):\n",
    "        target = torch.sparse.torch.eye(10).index_select(dim=0, index=target)\n",
    "        if(USE_CUDA):\n",
    "            data, target = data.cuda(), target.cuda()\n",
    "\n",
    "        optimizer.zero_grad()\n",
    "        output, reconstructions, masked = capsule_net(data)\n",
    "        loss = capsule_net.loss(data, output, target, reconstructions)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        correct = torch.sum(torch.argmax(masked, 1) == torch.argmax(target, 1))\n",
    "        train_loss = loss.item()\n",
    "        total_loss += train_loss\n",
    "        if batch_id % 100 == 0:\n",
    "            tqdm.write(\"Epoch: [{}/{}], Batch: [{}/{}], train accuracy: {:.6f}, loss: {:.6f}\".format(\n",
    "                epoch,\n",
    "                N_EPOCHS,\n",
    "                batch_id + 1,\n",
    "                n_batch,\n",
    "                correct / float(BATCH_SIZE),\n",
    "                train_loss / float(BATCH_SIZE)\n",
    "                ))\n",
    "    tqdm.write('Epoch: [{}/{}], train loss: {:.6f}'.format(epoch,N_EPOCHS,total_loss / len(train_loader.dataset)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Loading Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torchvision\n",
    "import torchvision.transforms as transforms\n",
    "\n",
    "transform = transforms.Compose(\n",
    "    [transforms.ToTensor(),\n",
    "     transforms.Normalize((0.1307,), (0.3081,))])\n",
    "\n",
    "trainset = torchvision.datasets.MNIST(root='./data/mnist', train=True,\n",
    "                                        download=True, transform=transform)\n",
    "trainloader = torch.utils.data.DataLoader(trainset, batch_size=BATCH_SIZE,\n",
    "                                          shuffle=True, num_workers=2)\n",
    "\n",
    "testset = torchvision.datasets.MNIST(root='./data/mnist', train=False,\n",
    "                                       download=True, transform=transform)\n",
    "testloader = torch.utils.data.DataLoader(testset, batch_size=BATCH_SIZE,\n",
    "                                         shuffle=False, num_workers=2)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 1/235 [00:01<04:35,  1.18s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [1/30], Batch: [1/235], train accuracy: 0.117188, loss: 0.003517\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 43%|████▎     | 101/235 [00:57<01:18,  1.70it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [1/30], Batch: [101/235], train accuracy: 0.628906, loss: 0.003080\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 86%|████████▌ | 201/235 [01:56<00:19,  1.70it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [1/30], Batch: [201/235], train accuracy: 0.753906, loss: 0.002670\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 235/235 [02:16<00:00,  1.73it/s]\n",
      "  0%|          | 0/235 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [1/30], train loss: 0.002966\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 1/235 [00:00<02:42,  1.44it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [2/30], Batch: [1/235], train accuracy: 0.699219, loss: 0.002398\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 43%|████▎     | 101/235 [00:59<01:19,  1.70it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [2/30], Batch: [101/235], train accuracy: 0.820312, loss: 0.001477\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 86%|████████▌ | 201/235 [01:58<00:20,  1.70it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [2/30], Batch: [201/235], train accuracy: 0.890625, loss: 0.001005\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 235/235 [02:18<00:00,  1.70it/s]\n",
      "  0%|          | 0/235 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [2/30], train loss: 0.001490\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 1/235 [00:00<02:42,  1.44it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [3/30], Batch: [1/235], train accuracy: 0.906250, loss: 0.000893\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 43%|████▎     | 101/235 [00:59<01:18,  1.70it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [3/30], Batch: [101/235], train accuracy: 0.906250, loss: 0.000708\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 86%|████████▌ | 201/235 [01:58<00:20,  1.70it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [3/30], Batch: [201/235], train accuracy: 0.945312, loss: 0.000540\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 235/235 [02:18<00:00,  1.70it/s]\n",
      "  0%|          | 0/235 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [3/30], train loss: 0.000635\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 1/235 [00:00<02:43,  1.44it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [4/30], Batch: [1/235], train accuracy: 0.949219, loss: 0.000414\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 43%|████▎     | 101/235 [00:59<01:18,  1.70it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [4/30], Batch: [101/235], train accuracy: 0.964844, loss: 0.000363\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 86%|████████▌ | 201/235 [01:58<00:20,  1.69it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [4/30], Batch: [201/235], train accuracy: 0.968750, loss: 0.000313\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 235/235 [02:18<00:00,  1.70it/s]\n",
      "  0%|          | 0/235 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [4/30], train loss: 0.000427\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 1/235 [00:00<02:41,  1.45it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [5/30], Batch: [1/235], train accuracy: 0.960938, loss: 0.000361\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 43%|████▎     | 101/235 [00:59<01:18,  1.70it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [5/30], Batch: [101/235], train accuracy: 0.976562, loss: 0.000261\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 86%|████████▌ | 201/235 [01:58<00:19,  1.70it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [5/30], Batch: [201/235], train accuracy: 0.972656, loss: 0.000330\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 235/235 [02:17<00:00,  1.70it/s]\n",
      "  0%|          | 0/235 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [5/30], train loss: 0.000355\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 1/235 [00:00<02:42,  1.44it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [6/30], Batch: [1/235], train accuracy: 0.957031, loss: 0.000341\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 43%|████▎     | 101/235 [00:59<01:18,  1.70it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [6/30], Batch: [101/235], train accuracy: 0.945312, loss: 0.000370\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 86%|████████▌ | 201/235 [01:58<00:20,  1.69it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [6/30], Batch: [201/235], train accuracy: 0.957031, loss: 0.000312\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 235/235 [02:18<00:00,  1.70it/s]\n",
      "  0%|          | 0/235 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [6/30], train loss: 0.000319\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 1/235 [00:00<02:44,  1.42it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [7/30], Batch: [1/235], train accuracy: 0.964844, loss: 0.000275\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 43%|████▎     | 101/235 [00:59<01:19,  1.69it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [7/30], Batch: [101/235], train accuracy: 0.964844, loss: 0.000298\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 86%|████████▌ | 201/235 [01:58<00:20,  1.70it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [7/30], Batch: [201/235], train accuracy: 0.976562, loss: 0.000213\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 235/235 [02:18<00:00,  1.70it/s]\n",
      "  0%|          | 0/235 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [7/30], train loss: 0.000274\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 1/235 [00:00<02:42,  1.44it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [8/30], Batch: [1/235], train accuracy: 0.964844, loss: 0.000259\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 43%|████▎     | 101/235 [00:59<01:18,  1.70it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [8/30], Batch: [101/235], train accuracy: 0.980469, loss: 0.000183\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 86%|████████▌ | 201/235 [01:58<00:20,  1.70it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [8/30], Batch: [201/235], train accuracy: 0.976562, loss: 0.000169\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 235/235 [02:18<00:00,  1.70it/s]\n",
      "  0%|          | 0/235 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [8/30], train loss: 0.000256\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 1/235 [00:00<02:42,  1.44it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [9/30], Batch: [1/235], train accuracy: 0.972656, loss: 0.000199\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 43%|████▎     | 101/235 [00:59<01:18,  1.70it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [9/30], Batch: [101/235], train accuracy: 0.976562, loss: 0.000265\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 86%|████████▌ | 201/235 [01:58<00:19,  1.70it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [9/30], Batch: [201/235], train accuracy: 0.957031, loss: 0.000324\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 235/235 [02:17<00:00,  1.70it/s]\n",
      "  0%|          | 0/235 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [9/30], train loss: 0.000229\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 1/235 [00:00<02:42,  1.44it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [10/30], Batch: [1/235], train accuracy: 0.980469, loss: 0.000177\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 43%|████▎     | 101/235 [00:59<01:18,  1.70it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [10/30], Batch: [101/235], train accuracy: 0.964844, loss: 0.000245\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 86%|████████▌ | 201/235 [01:58<00:20,  1.70it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [10/30], Batch: [201/235], train accuracy: 0.992188, loss: 0.000115\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 235/235 [02:17<00:00,  1.71it/s]\n",
      "  0%|          | 0/235 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [10/30], train loss: 0.000209\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 1/235 [00:00<02:41,  1.45it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [11/30], Batch: [1/235], train accuracy: 0.980469, loss: 0.000174\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 43%|████▎     | 101/235 [00:59<01:18,  1.70it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [11/30], Batch: [101/235], train accuracy: 0.972656, loss: 0.000198\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 86%|████████▌ | 201/235 [01:58<00:20,  1.70it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [11/30], Batch: [201/235], train accuracy: 0.957031, loss: 0.000221\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 235/235 [02:17<00:00,  1.71it/s]\n",
      "  0%|          | 0/235 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [11/30], train loss: 0.000201\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 1/235 [00:00<02:42,  1.44it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [12/30], Batch: [1/235], train accuracy: 0.984375, loss: 0.000164\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 43%|████▎     | 101/235 [00:59<01:18,  1.70it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [12/30], Batch: [101/235], train accuracy: 0.972656, loss: 0.000178\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 86%|████████▌ | 201/235 [01:58<00:19,  1.70it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [12/30], Batch: [201/235], train accuracy: 0.949219, loss: 0.000270\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 235/235 [02:17<00:00,  1.71it/s]\n",
      "  0%|          | 0/235 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [12/30], train loss: 0.000188\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 1/235 [00:00<02:42,  1.44it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [13/30], Batch: [1/235], train accuracy: 0.976562, loss: 0.000219\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 43%|████▎     | 101/235 [00:59<01:18,  1.70it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [13/30], Batch: [101/235], train accuracy: 0.968750, loss: 0.000251\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 86%|████████▌ | 201/235 [01:58<00:20,  1.70it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [13/30], Batch: [201/235], train accuracy: 0.988281, loss: 0.000142\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 235/235 [02:17<00:00,  1.71it/s]\n",
      "  0%|          | 0/235 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [13/30], train loss: 0.000178\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 1/235 [00:00<02:44,  1.43it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [14/30], Batch: [1/235], train accuracy: 0.980469, loss: 0.000177\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 43%|████▎     | 101/235 [00:59<01:18,  1.70it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [14/30], Batch: [101/235], train accuracy: 0.972656, loss: 0.000169\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 86%|████████▌ | 201/235 [01:57<00:19,  1.70it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [14/30], Batch: [201/235], train accuracy: 0.988281, loss: 0.000133\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 235/235 [02:17<00:00,  1.71it/s]\n",
      "  0%|          | 0/235 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [14/30], train loss: 0.000173\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 1/235 [00:00<02:44,  1.42it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [15/30], Batch: [1/235], train accuracy: 0.984375, loss: 0.000149\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 43%|████▎     | 101/235 [00:59<01:18,  1.70it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [15/30], Batch: [101/235], train accuracy: 0.976562, loss: 0.000152\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 86%|████████▌ | 201/235 [01:58<00:20,  1.69it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [15/30], Batch: [201/235], train accuracy: 0.988281, loss: 0.000113\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 235/235 [02:17<00:00,  1.70it/s]\n",
      "  0%|          | 0/235 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [15/30], train loss: 0.000164\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 1/235 [00:00<02:43,  1.43it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [16/30], Batch: [1/235], train accuracy: 0.976562, loss: 0.000148\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 43%|████▎     | 101/235 [00:59<01:18,  1.70it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [16/30], Batch: [101/235], train accuracy: 0.972656, loss: 0.000218\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 86%|████████▌ | 201/235 [01:58<00:19,  1.70it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [16/30], Batch: [201/235], train accuracy: 0.964844, loss: 0.000249\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 235/235 [02:17<00:00,  1.71it/s]\n",
      "  0%|          | 0/235 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [16/30], train loss: 0.000160\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 1/235 [00:00<02:42,  1.44it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [17/30], Batch: [1/235], train accuracy: 0.980469, loss: 0.000126\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 43%|████▎     | 101/235 [00:59<01:18,  1.71it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [17/30], Batch: [101/235], train accuracy: 0.976562, loss: 0.000159\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 86%|████████▌ | 201/235 [01:57<00:19,  1.71it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [17/30], Batch: [201/235], train accuracy: 0.968750, loss: 0.000247\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 235/235 [02:17<00:00,  1.71it/s]\n",
      "  0%|          | 0/235 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [17/30], train loss: 0.000152\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 1/235 [00:00<02:39,  1.46it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [18/30], Batch: [1/235], train accuracy: 0.988281, loss: 0.000106\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 43%|████▎     | 101/235 [00:59<01:18,  1.71it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [18/30], Batch: [101/235], train accuracy: 0.968750, loss: 0.000189\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 86%|████████▌ | 201/235 [01:57<00:19,  1.71it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [18/30], Batch: [201/235], train accuracy: 0.996094, loss: 0.000093\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 235/235 [02:17<00:00,  1.71it/s]\n",
      "  0%|          | 0/235 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [18/30], train loss: 0.000144\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 1/235 [00:00<02:40,  1.45it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [19/30], Batch: [1/235], train accuracy: 0.964844, loss: 0.000195\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 43%|████▎     | 101/235 [00:59<01:18,  1.71it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [19/30], Batch: [101/235], train accuracy: 0.980469, loss: 0.000130\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 86%|████████▌ | 201/235 [01:57<00:19,  1.71it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [19/30], Batch: [201/235], train accuracy: 0.980469, loss: 0.000130\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 235/235 [02:17<00:00,  1.71it/s]\n",
      "  0%|          | 0/235 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [19/30], train loss: 0.000154\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 1/235 [00:00<02:41,  1.45it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [20/30], Batch: [1/235], train accuracy: 0.988281, loss: 0.000089\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 43%|████▎     | 101/235 [00:59<01:18,  1.70it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [20/30], Batch: [101/235], train accuracy: 0.976562, loss: 0.000165\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 86%|████████▌ | 201/235 [01:58<00:20,  1.70it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [20/30], Batch: [201/235], train accuracy: 0.972656, loss: 0.000190\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 235/235 [02:17<00:00,  1.70it/s]\n",
      "  0%|          | 0/235 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [20/30], train loss: 0.000137\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 1/235 [00:00<02:41,  1.45it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [21/30], Batch: [1/235], train accuracy: 0.988281, loss: 0.000133\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 43%|████▎     | 101/235 [00:59<01:18,  1.70it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [21/30], Batch: [101/235], train accuracy: 0.976562, loss: 0.000142\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 86%|████████▌ | 201/235 [01:58<00:20,  1.70it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [21/30], Batch: [201/235], train accuracy: 0.988281, loss: 0.000124\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 235/235 [02:17<00:00,  1.70it/s]\n",
      "  0%|          | 0/235 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [21/30], train loss: 0.000131\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 1/235 [00:00<02:41,  1.45it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [22/30], Batch: [1/235], train accuracy: 0.968750, loss: 0.000192\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 43%|████▎     | 101/235 [00:59<01:19,  1.69it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [22/30], Batch: [101/235], train accuracy: 0.976562, loss: 0.000155\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 86%|████████▌ | 201/235 [01:58<00:20,  1.69it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [22/30], Batch: [201/235], train accuracy: 0.980469, loss: 0.000182\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 235/235 [02:18<00:00,  1.70it/s]\n",
      "  0%|          | 0/235 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [22/30], train loss: 0.000129\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 1/235 [00:00<02:43,  1.43it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [23/30], Batch: [1/235], train accuracy: 0.976562, loss: 0.000162\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 43%|████▎     | 101/235 [00:59<01:18,  1.70it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [23/30], Batch: [101/235], train accuracy: 0.957031, loss: 0.000224\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 86%|████████▌ | 201/235 [01:58<00:20,  1.70it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [23/30], Batch: [201/235], train accuracy: 0.972656, loss: 0.000174\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 235/235 [02:17<00:00,  1.70it/s]\n",
      "  0%|          | 0/235 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [23/30], train loss: 0.000126\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 1/235 [00:00<02:42,  1.44it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [24/30], Batch: [1/235], train accuracy: 0.980469, loss: 0.000145\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 43%|████▎     | 101/235 [00:59<01:18,  1.70it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [24/30], Batch: [101/235], train accuracy: 0.984375, loss: 0.000104\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 86%|████████▌ | 201/235 [01:58<00:19,  1.71it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [24/30], Batch: [201/235], train accuracy: 0.988281, loss: 0.000082\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 235/235 [02:17<00:00,  1.71it/s]\n",
      "  0%|          | 0/235 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [24/30], train loss: 0.000121\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 1/235 [00:00<02:41,  1.45it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [25/30], Batch: [1/235], train accuracy: 0.976562, loss: 0.000168\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 43%|████▎     | 101/235 [00:59<01:19,  1.69it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [25/30], Batch: [101/235], train accuracy: 1.000000, loss: 0.000055\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 86%|████████▌ | 201/235 [01:58<00:20,  1.70it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [25/30], Batch: [201/235], train accuracy: 0.980469, loss: 0.000190\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 235/235 [02:18<00:00,  1.70it/s]\n",
      "  0%|          | 0/235 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [25/30], train loss: 0.000121\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 1/235 [00:00<02:42,  1.44it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [26/30], Batch: [1/235], train accuracy: 0.972656, loss: 0.000151\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 43%|████▎     | 101/235 [00:59<01:19,  1.70it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [26/30], Batch: [101/235], train accuracy: 0.984375, loss: 0.000123\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 86%|████████▌ | 201/235 [01:58<00:20,  1.70it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [26/30], Batch: [201/235], train accuracy: 0.988281, loss: 0.000101\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 235/235 [02:18<00:00,  1.70it/s]\n",
      "  0%|          | 0/235 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [26/30], train loss: 0.000117\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 1/235 [00:00<02:43,  1.43it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [27/30], Batch: [1/235], train accuracy: 0.976562, loss: 0.000165\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 43%|████▎     | 101/235 [00:59<01:18,  1.70it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [27/30], Batch: [101/235], train accuracy: 0.980469, loss: 0.000148\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 86%|████████▌ | 201/235 [01:58<00:19,  1.70it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [27/30], Batch: [201/235], train accuracy: 0.988281, loss: 0.000092\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 235/235 [02:17<00:00,  1.70it/s]\n",
      "  0%|          | 0/235 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [27/30], train loss: 0.000113\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 1/235 [00:00<02:42,  1.44it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [28/30], Batch: [1/235], train accuracy: 0.968750, loss: 0.000180\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 43%|████▎     | 101/235 [00:59<01:18,  1.70it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [28/30], Batch: [101/235], train accuracy: 0.964844, loss: 0.000205\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 86%|████████▌ | 201/235 [01:58<00:20,  1.70it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [28/30], Batch: [201/235], train accuracy: 0.988281, loss: 0.000092\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 235/235 [02:17<00:00,  1.70it/s]\n",
      "  0%|          | 0/235 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [28/30], train loss: 0.000113\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 1/235 [00:00<02:42,  1.44it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [29/30], Batch: [1/235], train accuracy: 0.988281, loss: 0.000119\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 43%|████▎     | 101/235 [00:59<01:19,  1.69it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [29/30], Batch: [101/235], train accuracy: 0.980469, loss: 0.000140\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 86%|████████▌ | 201/235 [01:58<00:20,  1.69it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [29/30], Batch: [201/235], train accuracy: 0.968750, loss: 0.000168\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 235/235 [02:18<00:00,  1.70it/s]\n",
      "  0%|          | 0/235 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [29/30], train loss: 0.000111\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 1/235 [00:00<02:43,  1.43it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [30/30], Batch: [1/235], train accuracy: 0.988281, loss: 0.000117\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 43%|████▎     | 101/235 [00:59<01:19,  1.69it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [30/30], Batch: [101/235], train accuracy: 0.984375, loss: 0.000102\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 86%|████████▌ | 201/235 [01:58<00:19,  1.70it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [30/30], Batch: [201/235], train accuracy: 0.992188, loss: 0.000062\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 235/235 [02:18<00:00,  1.70it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [30/30], train loss: 0.000106\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "torch.manual_seed(1)\n",
    "\n",
    "\n",
    "config = Config()\n",
    "\n",
    "capsule_net = CapsNet(config)\n",
    "# capsule_net = torch.nn.DataParallel(capsule_net)\n",
    "if USE_CUDA:\n",
    "    capsule_net = capsule_net.cuda()\n",
    "\n",
    "optimizer = torch.optim.Adam(capsule_net.parameters())\n",
    "\n",
    "for e in range(1, N_EPOCHS + 1):\n",
    "    train(capsule_net, optimizer, trainloader, e)\n",
    "#     test(capsule_net, mnist.test_loader, e)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.save(capsule_net.state_dict(), \"./CapsNetMNIS.pth \")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "def test(capsule_net, test_loader, epoch):\n",
    "    capsule_net.eval()\n",
    "    test_loss = 0\n",
    "    correct = 0\n",
    "    for batch_id, (data, target) in enumerate(test_loader):\n",
    "\n",
    "        target = torch.sparse.torch.eye(10).index_select(dim=0, index=target)\n",
    "        if USE_CUDA:\n",
    "            data, target = data.cuda(), target.cuda()\n",
    "\n",
    "        output, reconstructions, masked = capsule_net(data)\n",
    "        loss = capsule_net.loss(data, output, target, reconstructions)\n",
    "\n",
    "        test_loss += loss.item()\n",
    "        correct += torch.sum(torch.argmax(masked, 1) == torch.argmax(target, 1))\n",
    "\n",
    "    tqdm.write(\n",
    "        \"Epoch: [{}/{}], test accuracy: {:.6f}, loss: {:.6f}\".format(epoch, N_EPOCHS, correct / len(test_loader.dataset),\n",
    "                                                                  test_loss / len(test_loader)))\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: [30/30], test accuracy: 0.985100, loss: 0.030051\n"
     ]
    }
   ],
   "source": [
    "test(capsule_net, testloader, 30)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Config for 16 1d vectors in Capsule Layer. Set the Softmax Dimension to 1 in this case\n",
    "# class Config:\n",
    "#     def __init__(self, dataset='mnist'):\n",
    "#         # CNN (cnn)\n",
    "#         self.cnn_in_channels = 1\n",
    "#         self.cnn_out_channels = 12\n",
    "#         self.cnn_kernel_size = 15\n",
    "\n",
    "#         # Primary Capsule (pc)\n",
    "#         self.pc_num_capsules = 1\n",
    "#         self.pc_in_channels = 12\n",
    "#         self.pc_out_channels = 16\n",
    "#         self.pc_kernel_size = 8\n",
    "#         self.pc_num_routes = 16 * 7 * 7\n",
    "\n",
    "#         # Digit Capsule 1 (dc)\n",
    "#         self.dc_num_capsules = 49\n",
    "#         self.dc_num_routes = 16 * 7 * 7\n",
    "#         self.dc_in_channels = 1\n",
    "#         self.dc_out_channels = 1 #16\n",
    "        \n",
    "#         # Digit Capsule 2 (dc)\n",
    "#         self.dc_2_num_capsules = 10\n",
    "#         self.dc_2_num_routes = 7 * 7\n",
    "#         self.dc_2_in_channels = 1 #16\n",
    "#         self.dc_2_out_channels = 16\n",
    "\n",
    "#         # Decoder\n",
    "#         self.input_width = 28\n",
    "#         self.input_height = 28"
   ]
  }
 ],
 "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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}