thesis_omniglot20w1s.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "import torch.nn.init as init\n",
    "import torch.nn.functional as F\n",
    "import torch.utils as utils\n",
    "import torch.utils.data\n",
    "import torchvision\n",
    "import torchvision.transforms as transforms\n",
    "import torchvision.transforms.functional as TF\n",
    "import torchvision.datasets as dsets\n",
    "from torchvision import models\n",
    "from torchvision.utils import save_image\n",
    "from torch.utils.tensorboard import SummaryWriter\n",
    "\n",
    "writer = SummaryWriter('runs/thesis06_20w_1s')\n",
    "\n",
    "\n",
    "from tqdm import tqdm\n",
    "\n",
    "from flows import PlanarFlow\n",
    "from utils import Binarize\n",
    "from codes import Linear_flipout, Flatten, count_parameters, EfficientNet\n",
    "\n",
    "from torchmeta.datasets import Omniglot, CIFARFS\n",
    "from torchmeta.transforms import Categorical, ClassSplitter, Rotation\n",
    "from torchvision.transforms import Compose, Resize, ToTensor\n",
    "from torchmeta.utils.data import BatchMetaDataLoader\n",
    "\n",
    "\n",
    "#from __future__ import print_function\n",
    "import argparse\n",
    "import cv2\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "import os\n",
    "cur_dir = \"C:/Users/KJH/OneDrive - skku.edu/KJH/Projects/2019winter_research\"\n",
    "#cur_dir = \"C:/Users/KJH-Laptop/OneDrive - skku.edu/KJH/Projects/2019winter_research/\"\n",
    "os.chdir(cur_dir)\n",
    "#os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"1\"\n",
    "os.environ[\"CUDA_LAUNCH_BLOCKING\"] = \"1\"\n",
    "\n",
    "import time\n",
    "import copy\n",
    "import random as rd\n",
    "\n",
    "device = torch.device('cuda')\n",
    "\n",
    "class net(nn.Module):\n",
    "    def __init__(self, num_classes):\n",
    "        super(net, self).__init__()\n",
    "        self.input_dim = [1, 28, 28]\n",
    "        self.num_classes = num_classes\n",
    "        \n",
    "        self.ctx = torch.hub.load('rwightman/gen-efficientnet-pytorch', 'efficientnet_b0', pretrained=True)\n",
    "        self.ctx.conv_stem = nn.Conv2d(1, 32, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
    "        self.ctx.classifier = nn.Sequential(\n",
    "            nn.BatchNorm1d(1280),\n",
    "            nn.Linear(1280, 320, bias = True),\n",
    "            nn.ReLU()\n",
    "        )\n",
    "        \n",
    "        self.layer = nn.LSTM(input_size=1280, hidden_size=1280, num_layers=5, bias=True, batch_first = True).to(device)\n",
    "        self.dec_mu = nn.ModuleList()\n",
    "        self.dec_logvar = nn.ModuleList()\n",
    "        \n",
    "        self.num_params = [[200 * 200, 200], [200 * 200, 200], [200 * self.num_classes, self.num_classes]]\n",
    "        \n",
    "        for layer_size in self.num_params:\n",
    "            self.dec_mu.append(\n",
    "                nn.Sequential(\n",
    "                    nn.ELU(),\n",
    "                    nn.Linear(1280, 1280, bias = True),\n",
    "                    nn.ELU(),\n",
    "                    nn.Linear(1280, layer_size[0] + layer_size[1], bias = True),\n",
    "                ).to(device))\n",
    "            \n",
    "            self.dec_logvar.append(\n",
    "                nn.Sequential(\n",
    "                    nn.ELU(),\n",
    "                    nn.Linear(1280, 1280, bias = True),\n",
    "                    nn.ELU(),\n",
    "                    nn.Linear(1280, layer_size[0] + layer_size[1], bias = True),\n",
    "                ).to(device))\n",
    "        \n",
    "        self.encoder = nn.Sequential(\n",
    "            nn.BatchNorm2d(1),\n",
    "            Flatten(),            \n",
    "            nn.Linear(784, 200, bias = False),\n",
    "            nn.ELU()\n",
    "        )\n",
    "\n",
    "            \n",
    "            \n",
    "    def forward(self, input_train, label_train, input_test, adapt_lr, adapt_step = 1):\n",
    "        ctx = self.ctx(input_train).view(input_train.shape[0], 5, -1)\n",
    "        h = torch.stack([ctx[torch.where(label_train == x)].mean(dim = 0) for x in range(self.num_classes)], dim = 1)\n",
    "        h = h.view(5, 1, -1)\n",
    "        \n",
    "        params = nn.ParameterList()\n",
    "        param_vals = []\n",
    "        kld = torch.tensor(0., device=device)\n",
    "        c = torch.empty([5, 1, 1280], requires_grad = False, device = device).fill_(0)\n",
    "        x_test_init = self.encoder(input_test)\n",
    "        x_test = x_test_init\n",
    "\n",
    "        for ind, (dec_mu, dec_logvar) in enumerate(zip(self.dec_mu, self.dec_logvar)):\n",
    "            x = torch.empty([1, 1, 1280], requires_grad=False, device = device).normal_(0, 1)\n",
    "            f, (h, c) = self.layer(x, (h, c))\n",
    "            \n",
    "            mu = dec_mu(f.view(1, -1)).squeeze()\n",
    "            logvar = dec_logvar(f.view(1, -1)).squeeze()\n",
    "            params.append(nn.Parameter(torch.stack((mu, logvar), dim = 0), requires_grad = True))\n",
    "        \n",
    "        optimizer = optim.SGD(params, lr = adapt_lr)\n",
    "        \n",
    "        for step in range(adapt_step + 1):\n",
    "            x_test = x_test_init\n",
    "            kld = torch.tensor(0., device=device)\n",
    "            for ind, param in enumerate(params):\n",
    "                param.retain_grad()\n",
    "\n",
    "                weight_mu = param[0, :self.num_params[ind][0]].view(-1, self.num_params[ind][1])\n",
    "                bias_mu = param[0, self.num_params[ind][0]:]\n",
    "\n",
    "                weight_logvar = param[1, :self.num_params[ind][0]].view(-1, self.num_params[ind][1])\n",
    "                bias_logvar = param[1, self.num_params[ind][0]:]\n",
    "\n",
    "                weight_noise = torch.empty(weight_mu.shape, requires_grad = False, device = device).normal_(0,1)\n",
    "                bias_noise = torch.empty(bias_mu.shape, requires_grad = False, device = device).normal_(0,1)\n",
    "                in_sign = torch.empty(x_test.shape, requires_grad = False, device = device).uniform_(-1,1).sign()\n",
    "                out_sign = torch.empty([x_test.shape[0], self.num_params[ind][1]], requires_grad = False, device = device).uniform_(-1,1).sign()\n",
    "\n",
    "                x_test = torch.mm(x_test, weight_mu) + torch.mm(in_sign * x_test, weight_noise * weight_mu * weight_logvar.div(2).exp()) * out_sign\n",
    "                x_test += (1 + bias_noise * bias_logvar.div(2).exp()) * bias_mu\n",
    "                x_test = F.elu(x_test)\n",
    "\n",
    "                kld += (mu.pow(2) - logvar + logvar.exp() - 1).mean()/2\n",
    "                \n",
    "            if step < adapt_step:\n",
    "                optimizer.zero_grad()\n",
    "                loss = F.cross_entropy(x_test, label_test) + 1e-6 * kld\n",
    "                loss.backward(retain_graph = True)\n",
    "                optimizer.step()\n",
    "\n",
    "        return x_test, kld"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using cache found in C:\\Users\\KJH/.cache\\torch\\hub\\rwightman_gen-efficientnet-pytorch_master\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  0) loss = 2.394786, kld = 0.000790, acc = 0.272292, time = 30.807 sec\n",
      " 10) loss = 1.950580, kld = 0.000167, acc = 0.444583, time = 31.331 sec\n",
      " 20) loss = 1.940853, kld = 0.000081, acc = 0.490208, time = 30.346 sec\n",
      " 30) loss = 1.970852, kld = 0.000058, acc = 0.460833, time = 30.190 sec\n",
      " 40) loss = 1.930782, kld = 0.000049, acc = 0.481042, time = 30.639 sec\n",
      " 50) loss = 1.805170, kld = 0.000046, acc = 0.506667, time = 31.029 sec\n",
      " 60) loss = 2.042447, kld = 0.000047, acc = 0.471042, time = 30.308 sec\n",
      " 70) loss = 1.923043, kld = 0.000044, acc = 0.497083, time = 29.834 sec\n",
      " 80) loss = 1.810271, kld = 0.000045, acc = 0.528958, time = 29.785 sec\n",
      " 90) loss = 2.008836, kld = 0.000044, acc = 0.478542, time = 29.577 sec\n",
      "100) loss = 1.655277, kld = 0.000041, acc = 0.536042, time = 30.207 sec\n",
      "110) loss = 1.872472, kld = 0.000042, acc = 0.516667, time = 31.147 sec\n",
      "120) loss = 1.817463, kld = 0.000040, acc = 0.508958, time = 31.313 sec\n",
      "130) loss = 1.812622, kld = 0.000039, acc = 0.535625, time = 31.167 sec\n",
      "140) loss = 1.987384, kld = 0.000039, acc = 0.498542, time = 30.837 sec\n",
      "150) loss = 1.844115, kld = 0.000037, acc = 0.509167, time = 30.031 sec\n",
      "160) loss = 1.928365, kld = 0.000039, acc = 0.498542, time = 30.555 sec\n",
      "170) loss = 1.788246, kld = 0.000038, acc = 0.525833, time = 29.824 sec\n",
      "180) loss = 1.832020, kld = 0.000036, acc = 0.501458, time = 30.202 sec\n",
      "190) loss = 1.732536, kld = 0.000037, acc = 0.536458, time = 30.283 sec\n"
     ]
    }
   ],
   "source": [
    "batch_size = 16\n",
    "meta_trainset = Omniglot('./data/',\n",
    "                   # Number of ways\n",
    "                   num_classes_per_task=20,\n",
    "                   # Resize the images to 28x28 and converts them to PyTorch tensors (from Torchvision)\n",
    "                   transform=Compose([Resize(28), ToTensor()]),\n",
    "                   # Transform the labels to integers (e.g. (\"Glagolitic/character01\", \"Sanskrit/character14\", ...) to (0, 1, ...))\n",
    "                   target_transform=Categorical(num_classes=20),\n",
    "                   # Creates new virtual classes with rotated versions of the images (from Santoro et al., 2016)\n",
    "                   class_augmentations=[Rotation([90, 180, 270])],\n",
    "                   meta_train=True,\n",
    "                   download=True)\n",
    "meta_trainset = ClassSplitter(meta_trainset, shuffle=True, num_train_per_class=1, num_test_per_class=15)\n",
    "meta_trainloader = BatchMetaDataLoader(meta_trainset, batch_size=batch_size, num_workers=0)\n",
    "\n",
    "model = net(20).cuda()\n",
    "optimizer = optim.Adam(model.parameters(), lr=1e-2)\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "\n",
    "num_batches = 200\n",
    "for batch_idx, meta_train_batch in zip(range(num_batches), meta_trainloader):\n",
    "    start = time.time()\n",
    "\n",
    "    train_inputs, train_targets = [x.to(device) for x in meta_train_batch[\"train\"]]\n",
    "    test_inputs, test_targets = [x.to(device) for x in meta_train_batch[\"test\"]]\n",
    "    \n",
    "    cum_loss = torch.tensor(0., device=device)\n",
    "    accuracy = torch.tensor(0., device=device)\n",
    "    reg = torch.tensor(0., device=device)\n",
    "\n",
    "    for task_idx, (train_input, train_target, test_input, test_target) in enumerate(\n",
    "        zip(train_inputs, train_targets, test_inputs, test_targets)):\n",
    "        optimizer.zero_grad()\n",
    "        pred, kld = model(train_input, train_target, test_input, test_target, 0.5, 5)\n",
    "        loss = criterion(pred, test_target)\n",
    "        (loss + 1e-6 * kld).backward()\n",
    "        optimizer.step()\n",
    "        with torch.no_grad():\n",
    "            cum_loss += loss\n",
    "            accuracy += torch.sum(pred.argmax(1) == test_target.cuda())\n",
    "            reg += kld\n",
    "    cum_loss /= batch_size\n",
    "    accuracy /= batch_size * 20 * 15\n",
    "    reg /= batch_size\n",
    "\n",
    "    writer.add_scalar('train/accuracy', accuracy, batch_idx)\n",
    "    writer.add_scalar('train/reg', reg, batch_idx)\n",
    "    if batch_idx % 10 == 0:\n",
    "        print(\"%3d) loss = %f, kld = %f, acc = %f, time = %.3f sec\" %(batch_idx, cum_loss, reg, accuracy, time.time() - start))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  0) loss = 2.215207, kld = 0.000035, acc = 0.457708, time = 30.405 sec\n",
      " 10) loss = 1.737480, kld = 0.000036, acc = 0.532500, time = 30.361 sec\n",
      " 20) loss = 1.829431, kld = 0.000036, acc = 0.526875, time = 30.432 sec\n",
      " 30) loss = 1.718449, kld = 0.000036, acc = 0.544583, time = 31.088 sec\n",
      " 40) loss = 1.893542, kld = 0.000036, acc = 0.513333, time = 30.800 sec\n",
      " 50) loss = 1.787285, kld = 0.000035, acc = 0.531250, time = 30.325 sec\n",
      " 60) loss = 1.624517, kld = 0.000035, acc = 0.544375, time = 29.889 sec\n",
      " 70) loss = 1.727175, kld = 0.000035, acc = 0.547917, time = 30.149 sec\n",
      " 80) loss = 1.728561, kld = 0.000036, acc = 0.545625, time = 30.365 sec\n",
      " 90) loss = 1.577610, kld = 0.000036, acc = 0.574792, time = 30.609 sec\n",
      "100) loss = 1.802424, kld = 0.000036, acc = 0.532708, time = 31.307 sec\n",
      "110) loss = 1.767332, kld = 0.000035, acc = 0.541875, time = 30.273 sec\n",
      "120) loss = 1.664457, kld = 0.000035, acc = 0.542500, time = 30.643 sec\n",
      "130) loss = 1.658135, kld = 0.000036, acc = 0.563333, time = 30.649 sec\n",
      "140) loss = 1.631656, kld = 0.000034, acc = 0.550625, time = 30.329 sec\n",
      "150) loss = 1.627812, kld = 0.000034, acc = 0.565833, time = 30.267 sec\n",
      "160) loss = 1.815804, kld = 0.000034, acc = 0.537083, time = 29.487 sec\n",
      "170) loss = 1.708923, kld = 0.000036, acc = 0.565625, time = 29.611 sec\n",
      "180) loss = 1.625499, kld = 0.000036, acc = 0.568750, time = 29.506 sec\n",
      "190) loss = 1.826722, kld = 0.000035, acc = 0.519583, time = 29.709 sec\n"
     ]
    }
   ],
   "source": [
    "optimizer = optim.Adam(model.parameters(), lr=1e-4)\n",
    "\n",
    "num_batches = 200\n",
    "for batch_idx, meta_train_batch in zip(range(num_batches), meta_trainloader):\n",
    "    start = time.time()\n",
    "\n",
    "    train_inputs, train_targets = [x.to(device) for x in meta_train_batch[\"train\"]]\n",
    "    test_inputs, test_targets = [x.to(device) for x in meta_train_batch[\"test\"]]\n",
    "    \n",
    "    cum_loss = torch.tensor(0., device=device)\n",
    "    accuracy = torch.tensor(0., device=device)\n",
    "    reg = torch.tensor(0., device=device)\n",
    "\n",
    "    for task_idx, (train_input, train_target, test_input, test_target) in enumerate(\n",
    "        zip(train_inputs, train_targets, test_inputs, test_targets)):\n",
    "        optimizer.zero_grad()\n",
    "        pred, kld = model(train_input, train_target, test_input, test_target, 0.5, 5)\n",
    "        loss = criterion(pred, test_target)\n",
    "        (loss + 1e-6 * kld).backward()\n",
    "        optimizer.step()\n",
    "        with torch.no_grad():\n",
    "            cum_loss += loss\n",
    "            accuracy += torch.sum(pred.argmax(1) == test_target.cuda())\n",
    "            reg += kld\n",
    "    cum_loss /= batch_size\n",
    "    accuracy /= batch_size * 20 * 15\n",
    "    reg /= batch_size\n",
    "\n",
    "    writer.add_scalar('train/accuracy', accuracy, batch_idx)\n",
    "    writer.add_scalar('train/reg', reg, batch_idx)\n",
    "    if batch_idx % 10 == 0:\n",
    "        print(\"%3d) loss = %f, kld = %f, acc = %f, time = %.3f sec\" %(batch_idx, cum_loss, reg, accuracy, time.time() - start))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  0) loss = 2.021474, kld = 0.000036, acc = 0.539375, time = 29.313 sec\n",
      " 10) loss = 1.719320, kld = 0.000035, acc = 0.546667, time = 30.608 sec\n",
      " 20) loss = 1.764859, kld = 0.000035, acc = 0.542500, time = 29.620 sec\n",
      " 30) loss = 1.694306, kld = 0.000036, acc = 0.549375, time = 29.098 sec\n",
      " 40) loss = 1.766228, kld = 0.000035, acc = 0.538125, time = 29.642 sec\n",
      " 50) loss = 1.758265, kld = 0.000035, acc = 0.553125, time = 29.412 sec\n",
      " 60) loss = 1.648570, kld = 0.000036, acc = 0.554167, time = 29.970 sec\n",
      " 70) loss = 1.867705, kld = 0.000037, acc = 0.537917, time = 29.251 sec\n",
      " 80) loss = 1.939738, kld = 0.000036, acc = 0.496042, time = 21.960 sec\n",
      " 90) loss = 1.781836, kld = 0.000034, acc = 0.522708, time = 16.538 sec\n",
      "100) loss = 1.736039, kld = 0.000036, acc = 0.564583, time = 16.906 sec\n",
      "110) loss = 1.806704, kld = 0.000035, acc = 0.513125, time = 16.668 sec\n",
      "120) loss = 1.845975, kld = 0.000035, acc = 0.519792, time = 16.440 sec\n",
      "130) loss = 1.954798, kld = 0.000036, acc = 0.504375, time = 16.461 sec\n",
      "140) loss = 1.547709, kld = 0.000036, acc = 0.592292, time = 29.166 sec\n",
      "150) loss = 1.673405, kld = 0.000035, acc = 0.555417, time = 29.483 sec\n",
      "160) loss = 1.753926, kld = 0.000035, acc = 0.535417, time = 29.085 sec\n",
      "170) loss = 1.608853, kld = 0.000035, acc = 0.572708, time = 29.400 sec\n",
      "180) loss = 1.696620, kld = 0.000036, acc = 0.534792, time = 29.578 sec\n",
      "190) loss = 1.826880, kld = 0.000036, acc = 0.513125, time = 29.733 sec\n"
     ]
    }
   ],
   "source": [
    "optimizer = optim.Adam(model.parameters(), lr=1e-6)\n",
    "\n",
    "num_batches = 200\n",
    "for batch_idx, meta_train_batch in zip(range(num_batches), meta_trainloader):\n",
    "    start = time.time()\n",
    "\n",
    "    train_inputs, train_targets = [x.to(device) for x in meta_train_batch[\"train\"]]\n",
    "    test_inputs, test_targets = [x.to(device) for x in meta_train_batch[\"test\"]]\n",
    "    \n",
    "    cum_loss = torch.tensor(0., device=device)\n",
    "    accuracy = torch.tensor(0., device=device)\n",
    "    reg = torch.tensor(0., device=device)\n",
    "\n",
    "    for task_idx, (train_input, train_target, test_input, test_target) in enumerate(\n",
    "        zip(train_inputs, train_targets, test_inputs, test_targets)):\n",
    "        optimizer.zero_grad()\n",
    "        pred, kld = model(train_input, train_target, test_input, test_target, 0.5, 5)\n",
    "        loss = criterion(pred, test_target)\n",
    "        (loss + 1e-6 * kld).backward()\n",
    "        optimizer.step()\n",
    "        with torch.no_grad():\n",
    "            cum_loss += loss\n",
    "            accuracy += torch.sum(pred.argmax(1) == test_target.cuda())\n",
    "            reg += kld\n",
    "    cum_loss /= batch_size\n",
    "    accuracy /= batch_size * 20 * 15\n",
    "    reg /= batch_size\n",
    "\n",
    "    writer.add_scalar('train/accuracy', accuracy, batch_idx)\n",
    "    writer.add_scalar('train/reg', reg, batch_idx)\n",
    "    if batch_idx % 10 == 0:\n",
    "        print(\"%3d) loss = %f, kld = %f, acc = %f, time = %.3f sec\" %(batch_idx, cum_loss, reg, accuracy, time.time() - start))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loss = 1.714735, test_kld = 0.000035, meta_test_acc = 0.547893, time = 12.386 sec\n"
     ]
    }
   ],
   "source": [
    "meta_testset  = Omniglot('./data/',\n",
    "                   # Number of ways\n",
    "                   num_classes_per_task=20,\n",
    "                   # Resize the images to 28x28 and converts them to PyTorch tensors (from Torchvision)\n",
    "                   transform=Compose([Resize(28), ToTensor()]),\n",
    "                   # Transform the labels to integers (e.g. (\"Glagolitic/character01\", \"Sanskrit/character14\", ...) to (0, 1, ...))\n",
    "                   target_transform=Categorical(num_classes=20),\n",
    "                   # Creates new virtual classes with rotated versions of the images (from Santoro et al., 2016)\n",
    "                   class_augmentations=[Rotation([90, 180, 270])],\n",
    "                   meta_test=True,\n",
    "                   download=True)\n",
    "meta_testset = ClassSplitter(meta_testset, shuffle=True, num_train_per_class=1, num_test_per_class=15)\n",
    "meta_testloader = BatchMetaDataLoader(meta_testset, batch_size=batch_size, num_workers=0)\n",
    "\n",
    "tot_loss = torch.tensor(0., device=device)\n",
    "tot_acc = torch.tensor(0., device=device)\n",
    "tot_reg = torch.tensor(0., device=device)\n",
    "\n",
    "for batch_idx, meta_test_batch in zip(range(num_batches), meta_testloader):\n",
    "    start = time.time()\n",
    "    train_inputs, train_targets = [x.to(device) for x in meta_test_batch[\"train\"]]\n",
    "    test_inputs, test_targets = [x.to(device) for x in meta_test_batch[\"test\"]]\n",
    "    \n",
    "    cum_loss = torch.tensor(0., device=device)\n",
    "    accuracy = torch.tensor(0., device=device)\n",
    "    reg = torch.tensor(0., device=device)\n",
    "\n",
    "    for task_idx, (train_input, train_target, test_input, test_target) in enumerate(\n",
    "        zip(train_inputs, train_targets, test_inputs, test_targets)):\n",
    "        optimizer.zero_grad()\n",
    "        pred, kld = model(train_input, train_target, test_input, test_target, 0.5, 5)\n",
    "        with torch.no_grad():            \n",
    "            loss = criterion(pred, test_target)\n",
    "            cum_loss += loss\n",
    "            accuracy += torch.sum(pred.argmax(1) == test_target.cuda())\n",
    "            reg += kld\n",
    "\n",
    "    tot_loss += cum_loss / batch_size\n",
    "    tot_acc += accuracy  / (batch_size * 20 * 15)\n",
    "    tot_reg += reg / batch_size\n",
    "        \n",
    "tot_loss /= num_batches\n",
    "tot_acc /= num_batches\n",
    "tot_reg /= num_batches\n",
    "        \n",
    "print(\"loss = %f, test_kld = %f, meta_test_acc = %f, time = %.3f sec\" %(tot_loss, tot_reg, tot_acc, time.time() - start))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.save(model.state_dict(), \"./save/thesis06_omniglot20w1s_5step\")"
   ]
  },
  {
   "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.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}