protonet layoutlm exp.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "42c06bd6-6242-4375-95ac-6cafb94e559a",
   "metadata": {},
   "source": [
    "# Protonet-LayoutLM with CORD\n",
    "\n",
    "References:\n",
    "\n",
    "https://github.com/cnielly/prototypical-networks-omniglot/blob/master/prototypical_networks_pytorch_omniglot.ipynb\n",
    "\n",
    "https://github.com/jakesnell/prototypical-networks/blob/master/protonets/models/few_shot.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "62b2a962-e3ea-48cf-81cb-a90dc5fac3a7",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'cuda'"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "import torch\n",
    "from torch import optim\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from PIL import Image\n",
    "from tqdm import tqdm_notebook\n",
    "from tqdm import trange\n",
    "from transformers import set_seed\n",
    "\n",
    "SEED = 42\n",
    "np.random.seed(SEED)\n",
    "torch.manual_seed(SEED)\n",
    "set_seed(SEED)\n",
    "\n",
    "DEVICE = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "DEVICE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "04fe6e1c-4503-467f-805a-6032f6543f1f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fe2f1e1b-50e6-4853-9e73-ec089c244e78",
   "metadata": {},
   "source": [
    "## Load data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "89afff23-ecc4-4e43-888b-86531d551357",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "dict_keys(['id', 'words', 'bboxes', 'ner_tags', 'image_path'])\n",
      "No. of labels:  23\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "022fc79b94f44b51b2139873b8164720",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "  0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "DatasetDict({\n",
      "    train: Dataset({\n",
      "        features: ['id', 'words', 'bboxes', 'ner_tags', 'image_path'],\n",
      "        num_rows: 800\n",
      "    })\n",
      "    test: Dataset({\n",
      "        features: ['id', 'words', 'bboxes', 'ner_tags', 'image_path'],\n",
      "        num_rows: 100\n",
      "    })\n",
      "    validation: Dataset({\n",
      "        features: ['id', 'words', 'bboxes', 'ner_tags', 'image_path'],\n",
      "        num_rows: 100\n",
      "    })\n",
      "})\n"
     ]
    }
   ],
   "source": [
    "from datasets import load_dataset_builder, load_dataset, concatenate_datasets\n",
    "from pprint import pprint\n",
    "import datasets\n",
    "\n",
    "datasets.utils.logging.set_verbosity_error()\n",
    "\n",
    "DATASET_URI = 'katanaml/cord'\n",
    "\n",
    "dsbuild = load_dataset_builder(DATASET_URI)\n",
    "pprint(dsbuild.info.features.keys())\n",
    "\n",
    "print(\"No. of labels: \", len(dsbuild.info.features['ner_tags'].feature.names))\n",
    "\n",
    "data = load_dataset(DATASET_URI)\n",
    "print(data)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3cda303c-039e-42e2-8099-98847c11e509",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Data sampling function\n",
    "\n",
    "Pick random sample of given support set size, query set size, no of classes and the dataset size. \n",
    "\n",
    "The function will be used to sample data every episode."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "58d4047b-fe43-4a89-a5d8-b83d37c06361",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "LABELS = dsbuild.info.features['ner_tags'].feature.names\n",
    "\n",
    "# select 15 for training, 8 for testing\n",
    "TRAINLABELS = np.random.choice(LABELS, 13, replace=False)\n",
    "TESTLABELS = np.random.choice(LABELS, 10, replace=False)\n",
    "\n",
    "HDIM = 768\n",
    "# TODO: remove 'O' label?"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "00200aa1-1e2a-4a36-b284-e2cf5c7ac169",
   "metadata": {},
   "source": [
    "Sampler below will be used for every episode to generate a sample of data to use in the episode.\n",
    "\n",
    "This will by default return all entities in a sample eventhough they are selected for a particular entity. We need additional preprocessing before the modeling stage to only train with the tokens that correspond to our target entity."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "2fa678fd-457f-4f8c-a683-e812c456187d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def sample_data(dataset, n_way, n_support, n_query, labels):\n",
    "    # randomly select B labels to sample in the\n",
    "    B = np.random.choice(labels, n_way, replace=False)\n",
    "    B = [LABELS.index(c) for c in B]\n",
    "    cls_sample = {}  # class -> sample\n",
    "    \n",
    "    # select support and query samples for each of these labels\n",
    "    for cls in B:\n",
    "        # select data containing the tag\n",
    "        sample = dataset.filter(lambda x: cls in x,\n",
    "                                input_columns=['ner_tags'],\n",
    "                               keep_in_memory=True)\n",
    "        # sample a small subset of the result\n",
    "        sample = sample.shuffle(keep_in_memory=True) \\\n",
    "                       .select(np.arange(n_support+n_query), \n",
    "                                keep_in_memory=True)\n",
    "        cls_sample[cls] = sample\n",
    "    \n",
    "    return {\n",
    "        'samples': cls_sample,\n",
    "        'n_way': n_way,\n",
    "        'n_support': n_support,\n",
    "        'n_query': n_query\n",
    "    }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "7e79af9b-c22e-46b9-984e-c03efed2a0c7",
   "metadata": {},
   "outputs": [],
   "source": [
    "id2label = data['test'].features['ner_tags'].feature.int2str\n",
    "\n",
    "def align_labels_with_tokens(labels, word_ids):\n",
    "    new_labels = []\n",
    "    current_word = None\n",
    "    for word_id in word_ids:\n",
    "        if word_id != current_word:\n",
    "            # Start of a new word!\n",
    "            current_word = word_id\n",
    "            label = -100 if word_id is None else id2label(labels[word_id])\n",
    "            new_labels.append(label)\n",
    "        elif word_id is None:\n",
    "            # Special token\n",
    "            new_labels.append(-100)\n",
    "        else:\n",
    "            # Same word as previous token\n",
    "            label = labels[word_id]\n",
    "            label = id2label(label)\n",
    "            new_labels.append(label)\n",
    "\n",
    "    return new_labels\n",
    "\n",
    "def align_all_labels(encoded, orig_data): \n",
    "    all_labels = []\n",
    "    for i in range(encoded['input_ids'].shape[0]): # for each instance\n",
    "        labels = align_labels_with_tokens(orig_data['ner_tags'][i], encoded.word_ids(i))\n",
    "        all_labels.append(labels)\n",
    "    return all_labels\n",
    "\n",
    "def getwords(wordids, words):\n",
    "    return [words[i]  if i is not None else -100 for i in wordids]\n",
    "\n",
    "def decode_tokens(input_ids):\n",
    "    decoded = []\n",
    "    for inp in input_ids:\n",
    "        decoded.append([processor.decode(i) for i in inp])\n",
    "    return decoded\n",
    "\n",
    "def load_image(data):\n",
    "    # update dataset with image\n",
    "    data['image'] = Image.open(data['image_path']).convert('RGB')\n",
    "    return data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "bfa3fa44-ce2e-433b-bff0-7f2b37109f85",
   "metadata": {},
   "outputs": [],
   "source": [
    "def encode(sample, processor, model):\n",
    "    enc = processor(sample['image'], sample['words'],\n",
    "                   boxes = sample['bboxes'],\n",
    "                    return_tensors='pt',\n",
    "                   truncation=True, padding=True)\n",
    "    with torch.no_grad():\n",
    "        enc = enc.to(DEVICE)\n",
    "        out = model(**enc)\n",
    "    out = out.last_hidden_state.detach()\n",
    "    out = out[:,:enc['input_ids'].shape[1],:]  # select only the word tokens\n",
    "    return enc, out\n",
    "\n",
    "def encode_and_filter(sample, processor, model, load_image_on_sample=False):\n",
    "    # output dimension: n_way, n_support+n_query, hiddendim\n",
    "    output_tensors = []\n",
    "    \n",
    "    for cls in sample['samples'].keys():\n",
    "        instance = sample['samples'][cls]\n",
    "        if load_image_on_sample:\n",
    "            instance = instance.map(load_image,\n",
    "                                    keep_in_memory=True,\n",
    "                                   num_proc=4\n",
    "                                   )\n",
    "        enc, out = encode(instance, processor, model)\n",
    "        aligned_labels = align_all_labels(enc, instance)\n",
    "        \n",
    "        # output dimension: n_support+n_query, hiddendim\n",
    "        selected = []\n",
    "        for i, tag in enumerate(aligned_labels):\n",
    "            targetclass = LABELS[cls]\n",
    "            # ids of tokens of the target class\n",
    "            idxs = [idx                \n",
    "                    for idx, tok in enumerate(tag)\n",
    "                    if tok == targetclass ]\n",
    "            # randomly select 1 token of the class\n",
    "            selidx = np.random.choice(idxs, 1)\n",
    "            selected.append(out[i,selidx,:])\n",
    "        selected = torch.cat(selected)\n",
    "        output_tensors.append(selected)\n",
    "        \n",
    "    output_tensors = torch.stack(output_tensors)\n",
    "    \n",
    "    assert output_tensors.shape == (sample['n_way'], sample['n_support']+sample['n_query'], HDIM)\n",
    "    return output_tensors"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3d2f9eac-cbd1-475c-a584-2a7fd9d39cd4",
   "metadata": {},
   "source": [
    "## Implement ProtoNet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "0993b0ac-4bc8-41c2-810e-a1c41f75e660",
   "metadata": {},
   "outputs": [],
   "source": [
    "class ProtoNet(nn.Module):\n",
    "    def __init__(self, encoder, processor, nnet):\n",
    "        super(ProtoNet, self).__init__()\n",
    "        self.encoder = encoder\n",
    "        self.processor = processor\n",
    "        self.nnet = nnet\n",
    "        \n",
    "        # exclude encoder params\n",
    "        for p in self.encoder.parameters():\n",
    "            p.requires_grad = False\n",
    "\n",
    "    def set_forward_loss(self, sample, load_image_on_sample=False):\n",
    "        n_way = sample['n_way']\n",
    "        n_support = sample['n_support']\n",
    "        n_query = sample['n_query']\n",
    "        \n",
    "        sample_enc = encode_and_filter(sample, self.processor, self.encoder, load_image_on_sample)\n",
    "\n",
    "        x_support = sample_enc[:, :n_support]\n",
    "        x_query = sample_enc[:, n_support:]\n",
    "\n",
    "        #target indices are 0 ... n_way-1\n",
    "        target_inds = torch.arange(0, n_way) \\\n",
    "                      .view(n_way, 1, 1)  \\\n",
    "                      .expand(n_way, n_query, 1).long()\n",
    "        target_inds = target_inds.to(DEVICE)\n",
    "\n",
    "        # encode images of the support and the query set\n",
    "        x = torch.cat([x_support.contiguous().view(n_way * n_support, *x_support.size()[2:]),\n",
    "                       x_query.contiguous().view(n_way * n_query, *x_query.size()[2:])], 0)\n",
    "\n",
    "        z = self.nnet.forward(x)\n",
    "        z_dim = z.size(-1)\n",
    "        z_proto = z[:n_way*n_support].view(n_way, n_support, z_dim).mean(1)\n",
    "        z_query = z[n_way*n_support:]\n",
    "\n",
    "        #compute distances\n",
    "        dists = self.euclidean_dist(z_query, z_proto)\n",
    "\n",
    "        #compute probabilities\n",
    "        log_p_y = F.log_softmax(-dists, dim=1).view(n_way, n_query, -1)\n",
    "\n",
    "        loss_val = -log_p_y.gather(2, target_inds).squeeze().view(-1).mean()\n",
    "        _, y_hat = log_p_y.max(2)\n",
    "        acc_val = torch.eq(y_hat, target_inds.squeeze()).float().mean()\n",
    "\n",
    "        return loss_val, {\n",
    "            'loss': loss_val.item(),\n",
    "            'acc': acc_val.item(),\n",
    "            'y_hat': y_hat\n",
    "            }\n",
    "    \n",
    "    def euclidean_dist(self, x, y):\n",
    "        # x: N x D\n",
    "        # y: M x D\n",
    "        n = x.size(0)\n",
    "        m = y.size(0)\n",
    "        d = x.size(1)\n",
    "        assert d == y.size(1)\n",
    "\n",
    "        x = x.unsqueeze(1).expand(n, m, d)\n",
    "        y = y.unsqueeze(0).expand(n, m, d)\n",
    "\n",
    "        return torch.pow(x - y, 2).sum(2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "03384fc3-d431-4a16-9f66-f72b9f6280af",
   "metadata": {},
   "source": [
    "## Training and Validation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "5949b62f-e07a-4f9b-afe5-062db9a290fd",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "def train(model, optimizer, scheduler,\n",
    "          dataset, n_way, n_support, n_query, \n",
    "          max_epoch, epoch_size,\n",
    "          load_image=False):\n",
    "    epoch = 0\n",
    "    \n",
    "    while epoch < max_epoch:\n",
    "        running_loss = 0\n",
    "        running_acc = 0\n",
    "        \n",
    "        for episode in trange(epoch_size, desc=\"Epoch {:d}/{:d} train\".format(epoch+1, max_epoch)):\n",
    "            sample = sample_data(dataset,\n",
    "                                 n_way, n_support, n_query,\n",
    "                                 TRAINLABELS)\n",
    "            optimizer.zero_grad()\n",
    "            loss, output = model.set_forward_loss(sample, load_image)\n",
    "            running_loss += output['loss']\n",
    "            running_acc += output['acc']\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "            \n",
    "        epoch_loss = running_loss / epoch_size\n",
    "        epoch_acc = running_acc / epoch_size\n",
    "        print('Epoch {:d} -- Loss: {:.4f} Acc: {:.4f}'.format(epoch+1,epoch_loss, epoch_acc))\n",
    "        epoch += 1\n",
    "        scheduler.step()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "dbcdb04a-a7a9-430e-979d-8ea922a518ed",
   "metadata": {},
   "outputs": [],
   "source": [
    "def test(model, \n",
    "         dataset, n_way, n_support, n_query, \n",
    "         test_episode,\n",
    "            load_image_on_sample=False):\n",
    "    running_loss = 0.0\n",
    "    running_acc = 0.0\n",
    "    for episode in trange(test_episode):\n",
    "        sample = sample_data(dataset,\n",
    "                             n_way, n_support, n_query, \n",
    "                             TESTLABELS)\n",
    "        loss, output = model.set_forward_loss(sample, load_image_on_sample)\n",
    "        running_loss += output['loss']\n",
    "        running_acc += output['acc']\n",
    "    avg_loss = running_loss / test_episode\n",
    "    avg_acc = running_acc / test_episode\n",
    "    print('Test results -- Loss: {:.4f} Acc: {:.4f}'.format(avg_loss, avg_acc))\n",
    "    return avg_acc"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4b4dc4a5-37d5-4a11-8c4c-43e24dc6b827",
   "metadata": {},
   "source": [
    "## Train the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "2aca1e47-6bc4-4117-a90d-cb57a435d577",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "N_WAY = 5\n",
    "N_SUPPORT = 5\n",
    "N_QUERY = 5\n",
    "\n",
    "MAX_EPOCHS = 10\n",
    "EPOCH_SIZE = 100"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "634e90ef-d644-4ec8-8856-300497f83255",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from transformers import AutoProcessor, AutoModelForTokenClassification, AutoModel\n",
    "processor = AutoProcessor.from_pretrained('microsoft/layoutlmv3-base', apply_ocr=False)\n",
    "model = AutoModel.from_pretrained('microsoft/layoutlmv3-base')\n",
    "model.eval()\n",
    "model.to(DEVICE);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "f66461ea-260f-405e-a630-5091ce33d553",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "nnet = nn.Sequential(\n",
    "    nn.Linear(HDIM, 256),\n",
    "    nn.Linear(256, 64),\n",
    ").to(DEVICE)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "bd523373-e879-427c-b771-3c849483c11f",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "pnet = ProtoNet(model, processor, nnet)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "e3f75b0d-e770-4dd0-94da-292f5bed6854",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "optimizer = optim.Adam(nnet.parameters(), lr=0.005)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "c961bc45-8cd0-40e7-85d2-c09b55100328",
   "metadata": {},
   "outputs": [],
   "source": [
    "scheduler = optim.lr_scheduler.StepLR(optimizer, 2, gamma=0.5, last_epoch=-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "99ec77b2-7cc5-430d-9bd2-3f1c22379914",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = concatenate_datasets([data['train'], data['test']])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "93ae1600-ec11-4130-929f-bd52e27992bb",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = dataset.map(load_image, keep_in_memory=True, num_proc=3)  \n",
    "load_image_on_sample = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "27f1df8c-ebc1-48e1-904e-4f62f1962cf6",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 1/2 train:   0%|          | 0/2 [00:00<?, ?it/s]"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Filter:   0%|          | 0/900 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Filter:   0%|          | 0/900 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Filter:   0%|          | 0/900 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Filter:   0%|          | 0/900 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Filter:   0%|          | 0/900 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map (num_proc=4):   0%|          | 0/10 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map (num_proc=4):   0%|          | 0/10 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map (num_proc=4):   0%|          | 0/10 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map (num_proc=4):   0%|          | 0/10 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map (num_proc=4):   0%|          | 0/10 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 1/2 train:  50%|█████     | 1/2 [00:19<00:19, 19.45s/it]"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Filter:   0%|          | 0/900 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Filter:   0%|          | 0/900 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Filter:   0%|          | 0/900 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Filter:   0%|          | 0/900 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Filter:   0%|          | 0/900 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map (num_proc=4):   0%|          | 0/10 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map (num_proc=4):   0%|          | 0/10 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map (num_proc=4):   0%|          | 0/10 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map (num_proc=4):   0%|          | 0/10 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map (num_proc=4):   0%|          | 0/10 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 1/2 train: 100%|██████████| 2/2 [00:36<00:00, 18.13s/it]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1 -- Loss: 4.4302 Acc: 0.5600\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 2/2 train:   0%|          | 0/2 [00:00<?, ?it/s]"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Filter:   0%|          | 0/900 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Filter:   0%|          | 0/900 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Filter:   0%|          | 0/900 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Filter:   0%|          | 0/900 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Filter:   0%|          | 0/900 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map (num_proc=4):   0%|          | 0/10 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map (num_proc=4):   0%|          | 0/10 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map (num_proc=4):   0%|          | 0/10 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map (num_proc=4):   0%|          | 0/10 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map (num_proc=4):   0%|          | 0/10 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 2/2 train:  50%|█████     | 1/2 [00:19<00:19, 19.06s/it]"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Filter:   0%|          | 0/900 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Filter:   0%|          | 0/900 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Filter:   0%|          | 0/900 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Filter:   0%|          | 0/900 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Filter:   0%|          | 0/900 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map (num_proc=4):   0%|          | 0/10 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map (num_proc=4):   0%|          | 0/10 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map (num_proc=4):   0%|          | 0/10 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map (num_proc=4):   0%|          | 0/10 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map (num_proc=4):   0%|          | 0/10 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch 2/2 train: 100%|██████████| 2/2 [00:40<00:00, 20.06s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 2 -- Loss: 8.5733 Acc: 0.4800\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "train(pnet, optimizer, scheduler,\n",
    "      dataset,\n",
    "      N_WAY, N_SUPPORT, N_QUERY,\n",
    "      MAX_EPOCHS, EPOCH_SIZE,\n",
    "      load_image_on_sample)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "333f698d-8084-42f5-b5af-a4d055cca18f",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 0/2 [00:00<?, ?it/s]"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Filter:   0%|          | 0/900 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Filter:   0%|          | 0/900 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Filter:   0%|          | 0/900 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Filter:   0%|          | 0/900 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Filter:   0%|          | 0/900 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map (num_proc=4):   0%|          | 0/10 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map (num_proc=4):   0%|          | 0/10 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map (num_proc=4):   0%|          | 0/10 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map (num_proc=4):   0%|          | 0/10 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map (num_proc=4):   0%|          | 0/10 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 50%|█████     | 1/2 [00:19<00:19, 19.26s/it]"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Filter:   0%|          | 0/900 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Filter:   0%|          | 0/900 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Filter:   0%|          | 0/900 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Filter:   0%|          | 0/900 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Filter:   0%|          | 0/900 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map (num_proc=4):   0%|          | 0/10 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map (num_proc=4):   0%|          | 0/10 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map (num_proc=4):   0%|          | 0/10 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map (num_proc=4):   0%|          | 0/10 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map (num_proc=4):   0%|          | 0/10 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 2/2 [00:38<00:00, 19.17s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test results -- Loss: 9.6178 Acc: 0.6400\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "acc = test(pnet,\n",
    "      dataset,\n",
    "      N_WAY, N_SUPPORT, N_QUERY,\n",
    "      EPOCH_SIZE,\n",
    "      load_image_on_sample)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "6f0fc530-c925-4ea1-a44b-1a625d205639",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from datetime import datetime"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "c0b83f76-91bc-4a5c-8803-a606cb4d180e",
   "metadata": {},
   "outputs": [],
   "source": [
    "OUTPATH = \"outputs\"\n",
    "filename = f\"pnet_{N_SUPPORT}_{N_QUERY}_{N_WAY}_{round(acc, 4)*100}_{datetime.now().strftime('%m-%d-%H:%M')}.pt\"\n",
    "torch.save(pnet, os.path.join(OUTPATH, filename))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.9.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}