thesis_omniglot5w1s.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, label_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_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 = 1.114917, kld = 0.000943, acc = 0.735833, time = 27.221 sec\n",
      " 10) loss = 0.800171, kld = 0.000491, acc = 0.778333, time = 26.762 sec\n",
      " 20) loss = 0.936648, kld = 0.000320, acc = 0.732500, time = 25.786 sec\n",
      " 30) loss = 0.770455, kld = 0.000244, acc = 0.827500, time = 26.028 sec\n",
      " 40) loss = 0.619009, kld = 0.000210, acc = 0.845833, time = 25.877 sec\n",
      " 50) loss = 0.615955, kld = 0.000189, acc = 0.830000, time = 26.454 sec\n",
      " 60) loss = 0.589267, kld = 0.000177, acc = 0.853333, time = 27.509 sec\n",
      " 70) loss = 0.745112, kld = 0.000165, acc = 0.789167, time = 25.629 sec\n",
      " 80) loss = 1.067363, kld = 0.000170, acc = 0.745833, time = 25.720 sec\n",
      " 90) loss = 0.511974, kld = 0.000170, acc = 0.851667, time = 26.368 sec\n",
      "100) loss = 0.877231, kld = 0.000165, acc = 0.837500, time = 26.684 sec\n",
      "110) loss = 0.622231, kld = 0.000160, acc = 0.837500, time = 27.811 sec\n",
      "120) loss = 0.454001, kld = 0.000169, acc = 0.861667, time = 25.911 sec\n",
      "130) loss = 0.338756, kld = 0.000164, acc = 0.897500, time = 25.784 sec\n",
      "140) loss = 0.481331, kld = 0.000162, acc = 0.869167, time = 26.099 sec\n",
      "150) loss = 0.836913, kld = 0.000159, acc = 0.849167, time = 25.902 sec\n",
      "160) loss = 0.452539, kld = 0.000158, acc = 0.866667, time = 27.734 sec\n",
      "170) loss = 0.579948, kld = 0.000160, acc = 0.871667, time = 26.916 sec\n",
      "180) loss = 0.360813, kld = 0.000155, acc = 0.918333, time = 25.632 sec\n",
      "190) loss = 0.379859, kld = 0.000155, acc = 0.889167, time = 25.388 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=1, 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, 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 * 75\n",
    "    reg /= batch_size\n",
    "\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.609192, kld = 0.000160, acc = 0.850000, time = 25.502 sec\n",
      " 10) loss = 1.584619, kld = 0.000159, acc = 0.868333, time = 25.921 sec\n",
      " 20) loss = 0.750604, kld = 0.000155, acc = 0.804167, time = 27.409 sec\n",
      " 30) loss = 0.798320, kld = 0.000159, acc = 0.827500, time = 25.369 sec\n",
      " 40) loss = 0.600944, kld = 0.000160, acc = 0.838333, time = 25.494 sec\n",
      " 50) loss = 0.345966, kld = 0.000160, acc = 0.906667, time = 25.701 sec\n",
      " 60) loss = 0.504387, kld = 0.000161, acc = 0.853333, time = 25.430 sec\n",
      " 70) loss = 0.309215, kld = 0.000161, acc = 0.902500, time = 27.540 sec\n",
      " 80) loss = 0.356428, kld = 0.000158, acc = 0.891667, time = 25.883 sec\n",
      " 90) loss = 0.472840, kld = 0.000162, acc = 0.875833, time = 25.415 sec\n",
      "100) loss = 0.558434, kld = 0.000162, acc = 0.871667, time = 25.283 sec\n",
      "110) loss = 0.192124, kld = 0.000157, acc = 0.949167, time = 25.623 sec\n",
      "120) loss = 0.934178, kld = 0.000159, acc = 0.915000, time = 26.508 sec\n",
      "130) loss = 0.412599, kld = 0.000157, acc = 0.874167, time = 26.836 sec\n",
      "140) loss = 0.416016, kld = 0.000157, acc = 0.894167, time = 25.857 sec\n",
      "150) loss = 0.624373, kld = 0.000161, acc = 0.822500, time = 25.692 sec\n",
      "160) loss = 0.441201, kld = 0.000151, acc = 0.880833, time = 26.306 sec\n",
      "170) loss = 0.275714, kld = 0.000165, acc = 0.910833, time = 26.421 sec\n",
      "180) loss = 0.623130, kld = 0.000156, acc = 0.866667, time = 28.040 sec\n",
      "190) loss = 1.300266, kld = 0.000158, acc = 0.761667, time = 26.540 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, 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 * 75\n",
    "    reg /= batch_size\n",
    "\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 = 0.624367, kld = 0.000154, acc = 0.831667, time = 26.549 sec\n",
      " 10) loss = 0.366524, kld = 0.000156, acc = 0.901667, time = 26.506 sec\n",
      " 20) loss = 0.432718, kld = 0.000163, acc = 0.883333, time = 25.655 sec\n",
      " 30) loss = 0.517042, kld = 0.000157, acc = 0.854167, time = 28.084 sec\n",
      " 40) loss = 0.727221, kld = 0.000157, acc = 0.869167, time = 26.317 sec\n",
      " 50) loss = 0.800768, kld = 0.000161, acc = 0.829167, time = 26.177 sec\n",
      " 60) loss = 0.392935, kld = 0.000156, acc = 0.895000, time = 25.389 sec\n",
      " 70) loss = 0.988804, kld = 0.000158, acc = 0.835000, time = 25.477 sec\n",
      " 80) loss = 0.379702, kld = 0.000154, acc = 0.905833, time = 27.630 sec\n",
      " 90) loss = 0.718591, kld = 0.000157, acc = 0.841667, time = 26.234 sec\n",
      "100) loss = 0.605882, kld = 0.000157, acc = 0.840000, time = 26.336 sec\n",
      "110) loss = 0.465253, kld = 0.000155, acc = 0.876667, time = 26.250 sec\n",
      "120) loss = 0.476437, kld = 0.000156, acc = 0.885833, time = 26.483 sec\n",
      "130) loss = 0.758112, kld = 0.000160, acc = 0.846667, time = 26.801 sec\n",
      "140) loss = 0.350693, kld = 0.000159, acc = 0.893333, time = 27.193 sec\n",
      "150) loss = 0.486982, kld = 0.000161, acc = 0.868333, time = 25.563 sec\n",
      "160) loss = 0.854985, kld = 0.000157, acc = 0.857500, time = 25.632 sec\n",
      "170) loss = 0.533162, kld = 0.000162, acc = 0.841667, time = 25.324 sec\n",
      "180) loss = 0.727609, kld = 0.000155, acc = 0.842500, time = 25.910 sec\n",
      "190) loss = 0.474154, kld = 0.000156, acc = 0.871667, time = 27.177 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, 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 * 75\n",
    "    reg /= batch_size\n",
    "\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": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  0) loss = 1.612343, kld = 0.000159, acc = 0.204167, time = 1.447 sec\n",
      "  1) loss = 1.614228, kld = 0.000161, acc = 0.206667, time = 1.021 sec\n",
      "  2) loss = 1.611155, kld = 0.000166, acc = 0.205000, time = 1.084 sec\n",
      "  3) loss = 1.617296, kld = 0.000155, acc = 0.180833, time = 1.039 sec\n",
      "  4) loss = 1.615215, kld = 0.000158, acc = 0.210833, time = 1.042 sec\n",
      "  5) loss = 1.613404, kld = 0.000156, acc = 0.204167, time = 1.017 sec\n",
      "  6) loss = 1.608475, kld = 0.000160, acc = 0.201667, time = 1.024 sec\n",
      "  7) loss = 1.611615, kld = 0.000155, acc = 0.188333, time = 1.030 sec\n",
      "  8) loss = 1.612790, kld = 0.000157, acc = 0.200000, time = 1.049 sec\n",
      "  9) loss = 1.613554, kld = 0.000157, acc = 0.198333, time = 1.066 sec\n",
      " 10) loss = 1.614281, kld = 0.000160, acc = 0.170833, time = 1.036 sec\n",
      " 11) loss = 1.614609, kld = 0.000160, acc = 0.190000, time = 1.055 sec\n",
      " 12) loss = 1.607516, kld = 0.000157, acc = 0.202500, time = 1.057 sec\n",
      " 13) loss = 1.613862, kld = 0.000158, acc = 0.214167, time = 1.082 sec\n",
      " 14) loss = 1.616034, kld = 0.000152, acc = 0.195833, time = 1.068 sec\n",
      " 15) loss = 1.610353, kld = 0.000153, acc = 0.215833, time = 1.031 sec\n",
      " 16) loss = 1.612136, kld = 0.000156, acc = 0.211667, time = 1.016 sec\n",
      " 17) loss = 1.612167, kld = 0.000161, acc = 0.217500, time = 1.049 sec\n",
      " 18) loss = 1.609811, kld = 0.000155, acc = 0.209167, time = 1.070 sec\n",
      " 19) loss = 1.613240, kld = 0.000157, acc = 0.199167, time = 1.037 sec\n",
      " 20) loss = 1.610410, kld = 0.000155, acc = 0.224167, time = 1.041 sec\n",
      " 21) loss = 1.613307, kld = 0.000153, acc = 0.195833, time = 1.101 sec\n",
      " 22) loss = 1.610273, kld = 0.000159, acc = 0.210000, time = 1.032 sec\n",
      " 23) loss = 1.612225, kld = 0.000153, acc = 0.200000, time = 1.065 sec\n",
      " 24) loss = 1.605365, kld = 0.000155, acc = 0.239167, time = 1.055 sec\n",
      " 25) loss = 1.613994, kld = 0.000153, acc = 0.194167, time = 1.062 sec\n",
      " 26) loss = 1.613543, kld = 0.000154, acc = 0.195000, time = 1.024 sec\n",
      " 27) loss = 1.609155, kld = 0.000162, acc = 0.204167, time = 1.042 sec\n",
      " 28) loss = 1.611467, kld = 0.000162, acc = 0.177500, time = 1.044 sec\n",
      " 29) loss = 1.614379, kld = 0.000155, acc = 0.176667, time = 1.082 sec\n",
      " 30) loss = 1.613627, kld = 0.000155, acc = 0.207500, time = 1.049 sec\n",
      " 31) loss = 1.612530, kld = 0.000163, acc = 0.204167, time = 1.062 sec\n",
      " 32) loss = 1.610655, kld = 0.000159, acc = 0.199167, time = 1.044 sec\n",
      " 33) loss = 1.615372, kld = 0.000156, acc = 0.175000, time = 1.045 sec\n",
      " 34) loss = 1.610494, kld = 0.000161, acc = 0.210000, time = 1.032 sec\n",
      " 35) loss = 1.615888, kld = 0.000160, acc = 0.188333, time = 1.080 sec\n",
      " 36) loss = 1.610367, kld = 0.000161, acc = 0.212500, time = 1.032 sec\n",
      " 37) loss = 1.606958, kld = 0.000153, acc = 0.205833, time = 1.020 sec\n",
      " 38) loss = 1.613576, kld = 0.000159, acc = 0.202500, time = 1.045 sec\n",
      " 39) loss = 1.607618, kld = 0.000157, acc = 0.225833, time = 1.034 sec\n",
      " 40) loss = 1.610383, kld = 0.000161, acc = 0.205000, time = 1.053 sec\n",
      " 41) loss = 1.610174, kld = 0.000159, acc = 0.226667, time = 1.030 sec\n",
      " 42) loss = 1.610553, kld = 0.000157, acc = 0.205833, time = 1.062 sec\n",
      " 43) loss = 1.612268, kld = 0.000166, acc = 0.200833, time = 1.037 sec\n",
      " 44) loss = 1.613756, kld = 0.000160, acc = 0.189167, time = 1.046 sec\n",
      " 45) loss = 1.608147, kld = 0.000158, acc = 0.225833, time = 1.039 sec\n",
      " 46) loss = 1.608814, kld = 0.000155, acc = 0.215833, time = 1.064 sec\n",
      " 47) loss = 1.614976, kld = 0.000159, acc = 0.200000, time = 1.054 sec\n",
      " 48) loss = 1.614697, kld = 0.000160, acc = 0.190000, time = 1.078 sec\n",
      " 49) loss = 1.612190, kld = 0.000158, acc = 0.205833, time = 1.114 sec\n"
     ]
    }
   ],
   "source": [
    "batch_size = 16\n",
    "num_batches = 200\n",
    "optimizer = optim.Adam(model.parameters(), lr=1e-7)\n",
    "criterion = nn.CrossEntropyLoss()\n",
    "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_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, 0)\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",
    "    print(\"%3d) loss = %f, kld = %f, acc = %f, time = %.3f sec\" %(batch_idx, cum_loss / batch_size, reg / batch_size, accuracy / (batch_size * 75), time.time() - start))\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": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.save(model.state_dict(), \"./save/thesis06_omniglot5w1s_5step\")"
   ]
  },
  {
   "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"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<All keys matched successfully>"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model = net(5).cuda()\n",
    "model.load_state_dict(torch.load(\"./save/thesis06_omniglot5w1s_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
}