thesis_omniglot5w5s_new.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_5w_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.Identity()\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], self.num_classes, 256)\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(self.num_classes, 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_train_init = self.encoder(input_train)\n",
    "        x_train = x_train_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):\n",
    "            x_train = x_train_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_train.shape, requires_grad = False, device = device).uniform_(-1,1).sign()\n",
    "                out_sign = torch.empty([x_train.shape[0], self.num_params[ind][1]], requires_grad = False, device = device).uniform_(-1,1).sign()\n",
    "\n",
    "                x_train = torch.mm(x_train, weight_mu) + torch.mm(in_sign * x_train, weight_noise * weight_mu * weight_logvar.div(2).exp()) * out_sign\n",
    "                x_train += (1 + bias_noise * bias_logvar.div(2).exp()) * bias_mu\n",
    "                x_train = F.elu(x_train)\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_train, label_train) + 1e-6 * kld\n",
    "                loss.backward(retain_graph = True)\n",
    "                optimizer.step()\n",
    "                \n",
    "                \n",
    "        x_test_init = self.encoder(input_test)\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",
    "        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 = 1.118820, kld = 0.000863, acc = 0.760000, time = 28.576 sec\n",
      " 10) loss = 0.986487, kld = 0.000395, acc = 0.692500, time = 28.577 sec\n",
      " 20) loss = 0.763360, kld = 0.000212, acc = 0.793333, time = 28.744 sec\n",
      " 30) loss = 0.499484, kld = 0.000131, acc = 0.840000, time = 28.278 sec\n",
      " 40) loss = 0.940768, kld = 0.000093, acc = 0.784167, time = 28.417 sec\n",
      " 50) loss = 0.510337, kld = 0.000075, acc = 0.870833, time = 27.220 sec\n",
      " 60) loss = 1.484791, kld = 0.000065, acc = 0.726667, time = 27.682 sec\n",
      " 70) loss = 0.431145, kld = 0.000057, acc = 0.858333, time = 27.924 sec\n",
      " 80) loss = 0.738859, kld = 0.000055, acc = 0.813333, time = 27.861 sec\n",
      " 90) loss = 0.621747, kld = 0.000053, acc = 0.803333, time = 28.043 sec\n",
      "100) loss = 0.665669, kld = 0.000052, acc = 0.798333, time = 27.888 sec\n",
      "110) loss = 0.600377, kld = 0.000051, acc = 0.845000, time = 27.758 sec\n",
      "120) loss = 0.502656, kld = 0.000051, acc = 0.860833, time = 28.028 sec\n",
      "130) loss = 1.088834, kld = 0.000049, acc = 0.778333, time = 28.349 sec\n",
      "140) loss = 0.596396, kld = 0.000048, acc = 0.828333, time = 28.397 sec\n",
      "150) loss = 0.462410, kld = 0.000049, acc = 0.879167, time = 27.796 sec\n",
      "160) loss = 0.384558, kld = 0.000049, acc = 0.876667, time = 28.598 sec\n",
      "170) loss = 0.633360, kld = 0.000050, acc = 0.860833, time = 28.511 sec\n",
      "180) loss = 0.828508, kld = 0.000049, acc = 0.765000, time = 15.483 sec\n",
      "190) loss = 0.626438, kld = 0.000049, acc = 0.833333, time = 15.206 sec\n"
     ]
    }
   ],
   "source": [
    "batch_size = 16\n",
    "meta_trainset = Omniglot('./data/',\n",
    "                   # Number of ways\n",
    "                   num_classes_per_task=5,\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=5),\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=5, num_test_per_class=15)\n",
    "meta_trainloader = BatchMetaDataLoader(meta_trainset, batch_size=batch_size, num_workers=0)\n",
    "\n",
    "model = net(5).cuda()\n",
    "optimizer = optim.Adam(model.parameters(), lr=1e-3)\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, 1, 1)\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 * 75\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 = 0.506735, kld = 0.000049, acc = 0.892500, time = 26.918 sec\n",
      " 10) loss = 0.407875, kld = 0.000047, acc = 0.880833, time = 27.400 sec\n",
      " 20) loss = 0.511529, kld = 0.000048, acc = 0.850833, time = 27.222 sec\n",
      " 30) loss = 0.850249, kld = 0.000047, acc = 0.794167, time = 27.018 sec\n",
      " 40) loss = 0.565616, kld = 0.000049, acc = 0.837500, time = 26.760 sec\n",
      " 50) loss = 0.417291, kld = 0.000048, acc = 0.880833, time = 26.750 sec\n",
      " 60) loss = 0.586819, kld = 0.000049, acc = 0.834167, time = 28.297 sec\n",
      " 70) loss = 0.483006, kld = 0.000047, acc = 0.850833, time = 26.863 sec\n",
      " 80) loss = 0.566484, kld = 0.000048, acc = 0.855833, time = 26.644 sec\n",
      " 90) loss = 0.285287, kld = 0.000048, acc = 0.929167, time = 27.441 sec\n",
      "100) loss = 0.838484, kld = 0.000050, acc = 0.856667, time = 26.387 sec\n",
      "110) loss = 0.455802, kld = 0.000047, acc = 0.850833, time = 27.261 sec\n",
      "120) loss = 1.181624, kld = 0.000050, acc = 0.815000, time = 27.434 sec\n",
      "130) loss = 0.543299, kld = 0.000049, acc = 0.831667, time = 27.518 sec\n",
      "140) loss = 0.375401, kld = 0.000049, acc = 0.897500, time = 27.816 sec\n",
      "150) loss = 0.704989, kld = 0.000048, acc = 0.814167, time = 27.465 sec\n",
      "160) loss = 0.509684, kld = 0.000048, acc = 0.875000, time = 26.584 sec\n",
      "170) loss = 0.384194, kld = 0.000050, acc = 0.881667, time = 27.786 sec\n",
      "180) loss = 0.593617, kld = 0.000049, acc = 0.854167, time = 26.338 sec\n",
      "190) loss = 0.717463, kld = 0.000047, acc = 0.857500, time = 26.562 sec\n"
     ]
    }
   ],
   "source": [
    "optimizer = optim.Adam(model.parameters(), lr=1e-5)\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, 1, 1)\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 * 75\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 = 1.147690, kld = 0.000048, acc = 0.868333, time = 28.237 sec\n",
      " 10) loss = 0.157103, kld = 0.000047, acc = 0.959167, time = 26.679 sec\n",
      " 20) loss = 0.371660, kld = 0.000048, acc = 0.888333, time = 26.710 sec\n",
      " 30) loss = 0.850984, kld = 0.000048, acc = 0.885833, time = 26.666 sec\n",
      " 40) loss = 0.423712, kld = 0.000048, acc = 0.912500, time = 26.874 sec\n",
      " 50) loss = 0.658808, kld = 0.000050, acc = 0.812500, time = 26.856 sec\n",
      " 60) loss = 0.478345, kld = 0.000049, acc = 0.840833, time = 27.026 sec\n",
      " 70) loss = 0.696860, kld = 0.000049, acc = 0.802500, time = 27.617 sec\n",
      " 80) loss = 0.639119, kld = 0.000047, acc = 0.855833, time = 27.302 sec\n",
      " 90) loss = 0.444986, kld = 0.000049, acc = 0.879167, time = 26.494 sec\n",
      "100) loss = 0.531235, kld = 0.000047, acc = 0.857500, time = 26.444 sec\n",
      "110) loss = 2.056088, kld = 0.000049, acc = 0.801667, time = 28.118 sec\n",
      "120) loss = 0.506234, kld = 0.000048, acc = 0.883333, time = 27.576 sec\n",
      "130) loss = 0.518983, kld = 0.000047, acc = 0.893333, time = 27.453 sec\n",
      "140) loss = 0.449749, kld = 0.000049, acc = 0.870833, time = 28.086 sec\n",
      "150) loss = 0.374231, kld = 0.000049, acc = 0.895000, time = 28.013 sec\n",
      "160) loss = 0.744665, kld = 0.000047, acc = 0.841667, time = 27.582 sec\n",
      "170) loss = 0.418599, kld = 0.000050, acc = 0.874167, time = 27.674 sec\n",
      "180) loss = 0.209768, kld = 0.000049, acc = 0.945000, time = 26.802 sec\n",
      "190) loss = 0.730047, kld = 0.000048, acc = 0.795000, time = 26.985 sec\n"
     ]
    }
   ],
   "source": [
    "optimizer = optim.Adam(model.parameters(), lr=1e-7)\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, 1, 1)\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 * 75\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": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loss = 0.597349, test_kld = 0.000049, meta_test_acc = 0.857650, time = 19.940 sec\n"
     ]
    }
   ],
   "source": [
    "meta_testset  = Omniglot('./data/',\n",
    "                   # Number of ways\n",
    "                   num_classes_per_task=5,\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=5),\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=False,\n",
    "                   download=True)\n",
    "meta_testset = ClassSplitter(meta_testset, shuffle=True, num_train_per_class=5, 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, 1, 1)\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 * 75)\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_omniglot5w5s_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
}