Initial Commit

README.md

@@ -1,3 +1,7 @@
----
-license: mit
----
+# FSL_Subspace
+
+Codes for work on few-shot learning on chest x-ray images ([paper](https://openreview.net/pdf?id=AF97JZpgPe)).
+
+Check our [website](https://few-shot-learning-on-chest-x-ray.github.io/Project-Page/) for a brief summary of the paper.
+
+tl;dr : We propose a computationally efficient few-shot learning method for diagnosing chest X-rays, which uses an ensemble of random subspaces and a novel loss function to create well-separated clusters of training data in discriminative subspaces. Our method is almost 1.8 times faster than the popular t-SVD method for subspace decomposition and yields promising results on large-scale CXR datasets.

dataloader/CIFAR_FS copy.py

@@ -0,0 +1,470 @@
+# Dataloader of Gidaris & Komodakis, CVPR 2018
+# Adapted from:
+# https://github.com/gidariss/FewShotWithoutForgetting/blob/master/dataloader.py
+from __future__ import print_function
+
+import os
+import os.path
+import numpy as np
+import random
+import pickle
+import json
+import math
+
+import torch
+import torch.utils.data as data
+import torchvision
+import torchvision.datasets as datasets
+import torchvision.transforms as transforms
+import torchnet as tnt
+
+import h5py
+
+import cv2
+from PIL import Image
+from PIL import ImageEnhance
+
+from pdb import set_trace as breakpoint
+
+from torchvision.transforms.transforms import ToPILImage
+
+
+# Set the appropriate paths of the datasets here.
+_CIFAR_FS_DATASET_DIR = './cifar/CIFAR-FS/'
+
+
+def buildLabelIndex(labels):
+    label2inds = {}
+    for idx, label in enumerate(labels):
+        if label not in label2inds:
+            label2inds[label] = []
+        label2inds[label].append(idx)
+
+    return label2inds
+
+
+def load_data(file):
+    try:
+        with open(file, 'rb') as fo:
+            data = pickle.load(fo)
+        return data
+    except:
+        with open(file, 'rb') as f:
+            u = pickle._Unpickler(f)
+            u.encoding = 'latin1'
+            data = u.load()
+        return data
+
+
+class CIFAR_FS(data.Dataset):
+    def __init__(self, phase='train', do_not_use_random_transf=False):
+
+        assert(phase == 'train' or phase == 'val' or phase ==
+               'test' or phase == 'trainval')
+        self.phase = phase
+        self.name = 'CIFAR_FS_' + phase
+
+        print('Loading CIFAR-FS dataset - phase {0}'.format(phase))
+        file_train_categories_train_phase = os.path.join(
+            _CIFAR_FS_DATASET_DIR,
+            'CIFAR_FS_train.pickle')
+        file_train_categories_val_phase = os.path.join(
+            _CIFAR_FS_DATASET_DIR,
+            'CIFAR_FS_train.pickle')
+        file_train_categories_test_phase = os.path.join(
+            _CIFAR_FS_DATASET_DIR,
+            'CIFAR_FS_train.pickle')
+        file_val_categories_val_phase = os.path.join(
+            _CIFAR_FS_DATASET_DIR,
+            'CIFAR_FS_val.pickle')
+        file_test_categories_test_phase = os.path.join(
+            _CIFAR_FS_DATASET_DIR,
+            'CIFAR_FS_test.pickle')
+
+        if self.phase == 'train':
+            # During training phase we only load the training phase images
+            # of the training categories (aka base categories).
+            data_train = load_data(file_train_categories_train_phase)
+            self.data = data_train['data']
+            self.labels = data_train['labels']
+
+            self.label2ind = buildLabelIndex(self.labels)
+            self.labelIds = sorted(self.label2ind.keys())
+            self.num_cats = len(self.labelIds)
+            self.labelIds_base = self.labelIds
+            self.num_cats_base = len(self.labelIds_base)
+        elif self.phase == 'trainval':
+            # During training phase we only load the training phase images
+            # of the training categories (aka base categories).
+            data_train = load_data(file_train_categories_train_phase)
+            self.data = data_train['data']
+            self.labels = data_train['labels']
+            data_base = load_data(file_train_categories_val_phase)
+            data_novel = load_data(file_val_categories_val_phase)
+            self.data = np.concatenate(
+                [self.data, data_novel['data']], axis=0)
+            self.data = np.concatenate(
+                [self.data, data_base['data']], axis=0)
+
+            self.labels = np.concatenate(
+                [self.labels, data_novel['labels']], axis=0)
+            self.labels = np.concatenate(
+                [self.labels, data_base['labels']], axis=0)
+
+            self.label2ind = buildLabelIndex(self.labels)
+            self.labelIds = sorted(self.label2ind.keys())
+            self.num_cats = len(self.labelIds)
+            self.labelIds_base = self.labelIds
+            self.num_cats_base = len(self.labelIds_base)
+        elif self.phase == 'val' or self.phase == 'test':
+            if self.phase == 'test':
+                # load data that will be used for evaluating the recognition
+                # accuracy of the base categories.
+                data_base = load_data(file_train_categories_test_phase)
+                # load data that will be use for evaluating the few-shot recogniton
+                # accuracy on the novel categories.
+                data_novel = load_data(file_test_categories_test_phase)
+            else:  # phase=='val'
+                # load data that will be used for evaluating the recognition
+                # accuracy of the base categories.
+                data_base = load_data(file_train_categories_val_phase)
+                # load data that will be use for evaluating the few-shot recogniton
+                # accuracy on the novel categories.
+                data_novel = load_data(file_val_categories_val_phase)
+
+            self.data = np.concatenate(
+                [data_base['data'], data_novel['data']], axis=0)
+            self.labels = data_base['labels'] + data_novel['labels']
+
+            self.label2ind = buildLabelIndex(self.labels)
+            self.labelIds = sorted(self.label2ind.keys())
+            self.num_cats = len(self.labelIds)
+
+            self.labelIds_base = buildLabelIndex(data_base['labels']).keys()
+            self.labelIds_novel = buildLabelIndex(data_novel['labels']).keys()
+            self.num_cats_base = len(self.labelIds_base)
+            self.num_cats_novel = len(self.labelIds_novel)
+            intersection = set(self.labelIds_base) & set(self.labelIds_novel)
+            assert(len(intersection) == 0)
+        else:
+            raise ValueError('Not valid phase {0}'.format(self.phase))
+
+        mean_pix = [x/255.0 for x in [129.37731888,
+                                      124.10583864, 112.47758569]]
+
+        std_pix = [x/255.0 for x in [68.20947949, 65.43124043, 70.45866994]]
+
+        normalize = transforms.Normalize(mean=mean_pix, std=std_pix)
+
+        if (self.phase == 'test' or self.phase == 'val') or (do_not_use_random_transf == True):
+
+            self.transform = transforms.Compose([
+                transforms.ToPILImage(),
+                # lambda x: np.asarray(x),
+                transforms.ToTensor(),
+                normalize
+            ])
+        else:
+            self.transform = transforms.Compose([
+                transforms.ToPILImage(),
+                transforms.RandomCrop(32, padding=4),
+                transforms.ColorJitter(
+                    brightness=0.4, contrast=0.4, saturation=0.4),
+                transforms.RandomHorizontalFlip(),
+                transforms.ToTensor(),
+                # lambda x: np.asarray(x),
+                normalize
+            ])
+
+    def __getitem__(self, index):
+        img, label = self.data[index], self.labels[index]
+        # doing this so that it is consistent with all other datasets
+        # to return a PIL Image
+
+        # img = Image.fromarray(img)
+        if self.transform is not None:
+            img = self.transform(img)
+        return img, label
+
+    def __len__(self):
+        return len(self.data)
+
+
+class FewShotDataloader():
+    def __init__(self,
+                 dataset,
+                 nKnovel=5,  # number of novel categories.
+                 nKbase=-1,  # number of base categories.
+                 # number of training examples per novel category.
+                 nExemplars=1,
+                 # number of test examples for all the novel categories.
+                 nTestNovel=15*5,
+                 # number of test examples for all the base categories.
+                 nTestBase=15*5,
+                 batch_size=1,  # number of training episodes per batch.
+                 num_workers=4,
+                 epoch_size=2000,  # number of batches per epoch.
+                 ):
+
+        self.dataset = dataset
+        self.phase = self.dataset.phase
+        max_possible_nKnovel = (self.dataset.num_cats_base if self.phase == 'train' or self.phase == 'trainval'
+                                else self.dataset.num_cats_novel)
+        assert(nKnovel >= 0 and nKnovel < max_possible_nKnovel)
+        self.nKnovel = nKnovel
+
+        max_possible_nKbase = self.dataset.num_cats_base
+        nKbase = nKbase if nKbase >= 0 else max_possible_nKbase
+        if (self.phase == 'train' or self.phase == 'trainval') and nKbase > 0:
+            nKbase -= self.nKnovel
+            max_possible_nKbase -= self.nKnovel
+
+        assert(nKbase >= 0 and nKbase <= max_possible_nKbase)
+        self.nKbase = nKbase
+
+        self.nExemplars = nExemplars
+        self.nTestNovel = nTestNovel
+        self.nTestBase = nTestBase
+        self.batch_size = batch_size
+        self.epoch_size = epoch_size
+        self.num_workers = num_workers
+        self.is_eval_mode = (self.phase == 'test') or (self.phase == 'val')
+
+    def sampleImageIdsFrom(self, cat_id, sample_size=1):
+        """
+        Samples `sample_size` number of unique image ids picked from the
+        category `cat_id` (i.e., self.dataset.label2ind[cat_id]).
+
+        Args:
+            cat_id: a scalar with the id of the category from which images will
+                be sampled.
+            sample_size: number of images that will be sampled.
+
+        Returns:
+            image_ids: a list of length `sample_size` with unique image ids.
+        """
+        assert(cat_id in self.dataset.label2ind)
+        assert(len(self.dataset.label2ind[cat_id]) >= sample_size)
+        # Note: random.sample samples elements without replacement.
+        # seed = random.randint(1,10000000)
+        # random.seed(seed)
+        return random.sample(self.dataset.label2ind[cat_id], sample_size)
+
+    def sampleCategories(self, cat_set, sample_size=1):
+        """
+        Samples `sample_size` number of unique categories picked from the
+        `cat_set` set of categories. `cat_set` can be either 'base' or 'novel'.
+
+        Args:
+            cat_set: string that specifies the set of categories from which
+                categories will be sampled.
+            sample_size: number of categories that will be sampled.
+
+        Returns:
+            cat_ids: a list of length `sample_size` with unique category ids.
+        """
+        if cat_set == 'base':
+            labelIds = self.dataset.labelIds_base
+        elif cat_set == 'novel':
+            labelIds = self.dataset.labelIds_novel
+        else:
+            raise ValueError('Not recognized category set {}'.format(cat_set))
+
+        assert(len(labelIds) >= sample_size)
+        # return sample_size unique categories chosen from labelIds set of
+        # categories (that can be either self.labelIds_base or self.labelIds_novel)
+        # Note: random.sample samples elements without replacement.
+        return random.sample(labelIds, sample_size)
+
+    def sample_base_and_novel_categories(self, nKbase, nKnovel):
+        """
+        Samples `nKbase` number of base categories and `nKnovel` number of novel
+        categories.
+
+        Args:
+            nKbase: number of base categories
+            nKnovel: number of novel categories
+
+        Returns:
+            Kbase: a list of length 'nKbase' with the ids of the sampled base
+                categories.
+            Knovel: a list of lenght 'nKnovel' with the ids of the sampled novel
+                categories.
+        """
+        if self.is_eval_mode:
+            assert(nKnovel <= self.dataset.num_cats_novel)
+            # sample from the set of base categories 'nKbase' number of base
+            # categories.
+            Kbase = sorted(self.sampleCategories('base', nKbase))
+            # sample from the set of novel categories 'nKnovel' number of novel
+            # categories.
+            Knovel = sorted(self.sampleCategories('novel', nKnovel))
+        else:
+            # sample from the set of base categories 'nKnovel' + 'nKbase' number
+            # of categories.
+            cats_ids = self.sampleCategories('base', nKnovel+nKbase)
+            assert(len(cats_ids) == (nKnovel+nKbase))
+            # Randomly pick 'nKnovel' number of fake novel categories and keep
+            # the rest as base categories.
+            random.shuffle(cats_ids)
+            Knovel = sorted(cats_ids[:nKnovel])
+            Kbase = sorted(cats_ids[nKnovel:])
+
+        return Kbase, Knovel
+
+    def sample_test_examples_for_base_categories(self, Kbase, nTestBase):
+        """
+        Sample `nTestBase` number of images from the `Kbase` categories.
+
+        Args:
+            Kbase: a list of length `nKbase` with the ids of the categories from
+                where the images will be sampled.
+            nTestBase: the total number of images that will be sampled.
+
+        Returns:
+            Tbase: a list of length `nTestBase` with 2-element tuples. The 1st
+                element of each tuple is the image id that was sampled and the
+                2nd elemend is its category label (which is in the range
+                [0, len(Kbase)-1]).
+        """
+        Tbase = []
+        if len(Kbase) > 0:
+            # Sample for each base category a number images such that the total
+            # number sampled images of all categories to be equal to `nTestBase`.
+            KbaseIndices = np.random.choice(
+                np.arange(len(Kbase)), size=nTestBase, replace=True)
+            KbaseIndices, NumImagesPerCategory = np.unique(
+                KbaseIndices, return_counts=True)
+
+            for Kbase_idx, NumImages in zip(KbaseIndices, NumImagesPerCategory):
+                imd_ids = self.sampleImageIdsFrom(
+                    Kbase[Kbase_idx], sample_size=NumImages)
+                Tbase += [(img_id, Kbase_idx) for img_id in imd_ids]
+
+        assert(len(Tbase) == nTestBase)
+
+        return Tbase
+
+    def sample_train_and_test_examples_for_novel_categories(
+            self, Knovel, nTestNovel, nExemplars, nKbase):
+        """Samples train and test examples of the novel categories.
+
+        Args:
+            Knovel: a list with the ids of the novel categories.
+            nTestNovel: the total number of test images that will be sampled
+                from all the novel categories.
+            nExemplars: the number of training examples per novel category that
+                will be sampled.
+            nKbase: the number of base categories. It is used as offset of the
+                category index of each sampled image.
+
+        Returns:
+            Tnovel: a list of length `nTestNovel` with 2-element tuples. The
+                1st element of each tuple is the image id that was sampled and
+                the 2nd element is its category label (which is in the range
+                [nKbase, nKbase + len(Knovel) - 1]).
+            Exemplars: a list of length len(Knovel) * nExemplars of 2-element
+                tuples. The 1st element of each tuple is the image id that was
+                sampled and the 2nd element is its category label (which is in
+                the ragne [nKbase, nKbase + len(Knovel) - 1]).
+        """
+
+        if len(Knovel) == 0:
+            return [], []
+
+        nKnovel = len(Knovel)
+        Tnovel = []
+        Exemplars = []
+        assert((nTestNovel % nKnovel) == 0)
+        nEvalExamplesPerClass = int(nTestNovel / nKnovel)
+
+        for Knovel_idx in range(len(Knovel)):
+            imd_ids = self.sampleImageIdsFrom(
+                Knovel[Knovel_idx],
+                sample_size=(nEvalExamplesPerClass + nExemplars))
+
+            imds_tnovel = imd_ids[:nEvalExamplesPerClass]
+            imds_ememplars = imd_ids[nEvalExamplesPerClass:]
+
+            Tnovel += [(img_id, nKbase+Knovel_idx) for img_id in imds_tnovel]
+            Exemplars += [(img_id, nKbase+Knovel_idx)
+                          for img_id in imds_ememplars]
+        assert(len(Tnovel) == nTestNovel)
+        assert(len(Exemplars) == len(Knovel) * nExemplars)
+        random.shuffle(Exemplars)
+
+        return Tnovel, Exemplars
+
+    def sample_episode(self):
+        """Samples a training episode."""
+        nKnovel = self.nKnovel
+        nKbase = self.nKbase
+        nTestNovel = self.nTestNovel
+        nTestBase = self.nTestBase
+        nExemplars = self.nExemplars
+
+        Kbase, Knovel = self.sample_base_and_novel_categories(nKbase, nKnovel)
+        Tbase = self.sample_test_examples_for_base_categories(Kbase, nTestBase)
+        Tnovel, Exemplars = self.sample_train_and_test_examples_for_novel_categories(
+            Knovel, nTestNovel, nExemplars, nKbase)
+
+        # concatenate the base and novel category examples.
+        Test = Tbase + Tnovel
+        random.shuffle(Test)
+        Kall = Kbase + Knovel
+
+        return Exemplars, Test, Kall, nKbase
+
+    def createExamplesTensorData(self, examples):
+        """
+        Creates the examples image and label tensor data.
+
+        Args:
+            examples: a list of 2-element tuples, each representing a
+                train or test example. The 1st element of each tuple
+                is the image id of the example and 2nd element is the
+                category label of the example, which is in the range
+                [0, nK - 1], where nK is the total number of categories
+                (both novel and base).
+
+        Returns:
+            images: a tensor of shape [nExamples, Height, Width, 3] with the
+                example images, where nExamples is the number of examples
+                (i.e., nExamples = len(examples)).
+            labels: a tensor of shape [nExamples] with the category label
+                of each example.
+        """
+        images = torch.stack(
+            [self.dataset[img_idx][0] for img_idx, _ in examples], dim=0)
+        labels = torch.LongTensor([label for _, label in examples])
+        return images, labels
+
+    def get_iterator(self, epoch=0):
+        rand_seed = epoch
+        random.seed(rand_seed)
+        np.random.seed(rand_seed)
+
+        def load_function(iter_idx):
+            Exemplars, Test, Kall, nKbase = self.sample_episode()
+            Xt, Yt = self.createExamplesTensorData(Test)
+            Kall = torch.LongTensor(Kall)
+            if len(Exemplars) > 0:
+                Xe, Ye = self.createExamplesTensorData(Exemplars)
+                return Xe, Ye, Xt, Yt, Kall, nKbase
+            else:
+                return Xt, Yt, Kall, nKbase
+
+        tnt_dataset = tnt.dataset.ListDataset(
+            elem_list=range(self.epoch_size), load=load_function)
+        data_loader = tnt_dataset.parallel(
+            batch_size=self.batch_size,
+            num_workers=(0 if self.is_eval_mode else self.num_workers),
+            shuffle=(False if self.is_eval_mode else True))
+
+        return data_loader
+
+    def __call__(self, epoch=0):
+        return self.get_iterator(epoch)
+
+    def __len__(self):
+        return int(self.epoch_size / self.batch_size)

dataloader/CIFAR_FS.py

@@ -0,0 +1,488 @@
+# Dataloader of Gidaris & Komodakis, CVPR 2018
+# Adapted from:
+# https://github.com/gidariss/FewShotWithoutForgetting/blob/master/dataloader.py
+from __future__ import print_function
+
+import os
+import os.path
+import numpy as np
+import random
+import pickle
+import json
+import math
+
+import torch
+import torch.utils.data as data
+import torchvision
+import torchvision.datasets as datasets
+import torchvision.transforms as transforms
+import torchnet as tnt
+
+import h5py
+
+import cv2
+from PIL import Image
+from PIL import ImageEnhance
+import matplotlib.pyplot as plt
+
+from pdb import set_trace as breakpoint
+
+from torchvision.transforms.transforms import ToPILImage
+
+
+# Set the appropriate paths of the datasets here.
+_CIFAR_FS_DATASET_DIR = './cifar/CIFAR-FS/'
+
+
+def buildLabelIndex(labels):
+    label2inds = {}
+    for idx, label in enumerate(labels):
+        if label not in label2inds:
+            label2inds[label] = []
+        label2inds[label].append(idx)
+
+    return label2inds
+
+
+def load_data(file):
+    try:
+        with open(file, 'rb') as fo:
+            data = pickle.load(fo)
+        return data
+    except:
+        with open(file, 'rb') as f:
+            u = pickle._Unpickler(f)
+            u.encoding = 'latin1'
+            data = u.load()
+        return data
+
+
+class CIFAR_FS(data.Dataset):
+    def __init__(self, phase='train', do_not_use_random_transf=False):
+
+        assert(phase == 'train' or phase == 'val' or phase ==
+               'test' or phase == 'trainval')
+        self.phase = phase
+        self.name = 'CIFAR_FS_' + phase
+
+        print('Loading CIFAR-FS dataset - phase {0}'.format(phase))
+        file_train_categories_train_phase = os.path.join(
+            _CIFAR_FS_DATASET_DIR,
+            'CIFAR_FS_train.pickle')
+        file_train_categories_val_phase = os.path.join(
+            _CIFAR_FS_DATASET_DIR,
+            'CIFAR_FS_train.pickle')
+        file_train_categories_test_phase = os.path.join(
+            _CIFAR_FS_DATASET_DIR,
+            'CIFAR_FS_train.pickle')
+        file_val_categories_val_phase = os.path.join(
+            _CIFAR_FS_DATASET_DIR,
+            'CIFAR_FS_val.pickle')
+        file_test_categories_test_phase = os.path.join(
+            _CIFAR_FS_DATASET_DIR,
+            'CIFAR_FS_test.pickle')
+
+        if self.phase == 'train':
+            # During training phase we only load the training phase images
+            # of the training categories (aka base categories).
+            data_train = load_data(file_train_categories_train_phase)
+            self.data = data_train['data']
+            self.labels = data_train['labels']
+
+            self.label2ind = buildLabelIndex(self.labels)
+            self.labelIds = sorted(self.label2ind.keys())
+
+            self.num_cats = len(self.labelIds)
+            self.labelIds_base = self.labelIds
+            self.num_cats_base = len(self.labelIds_base)
+        elif self.phase == 'trainval':
+            # During training phase we only load the training phase images
+            # of the training categories (aka base categories).
+            data_train = load_data(file_train_categories_train_phase)
+            self.data = data_train['data']
+            self.labels = data_train['labels']
+            data_base = load_data(file_train_categories_val_phase)
+            data_novel = load_data(file_val_categories_val_phase)
+            self.data = np.concatenate(
+                [self.data, data_novel['data']], axis=0)
+            self.data = np.concatenate(
+                [self.data, data_base['data']], axis=0)
+
+            self.labels = np.concatenate(
+                [self.labels, data_novel['labels']], axis=0)
+            self.labels = np.concatenate(
+                [self.labels, data_base['labels']], axis=0)
+
+            self.label2ind = buildLabelIndex(self.labels)
+            self.labelIds = sorted(self.label2ind.keys())
+            self.num_cats = len(self.labelIds)
+            self.labelIds_base = self.labelIds
+            self.num_cats_base = len(self.labelIds_base)
+        elif self.phase == 'val' or self.phase == 'test':
+            if self.phase == 'test':
+                # load data that will be used for evaluating the recognition
+                # accuracy of the base categories.
+                data_base = load_data(file_train_categories_test_phase)
+                # load data that will be use for evaluating the few-shot recogniton
+                # accuracy on the novel categories.
+                data_novel = load_data(file_test_categories_test_phase)
+            else:  # phase=='val'
+                # load data that will be used for evaluating the recognition
+                # accuracy of the base categories.
+                data_base = load_data(file_train_categories_val_phase)
+                # load data that will be use for evaluating the few-shot recogniton
+                # accuracy on the novel categories.
+                data_novel = load_data(file_val_categories_val_phase)
+
+            self.data = np.concatenate(
+                [data_base['data'], data_novel['data']], axis=0)
+            self.labels = data_base['labels'] + data_novel['labels']
+
+            self.label2ind = buildLabelIndex(self.labels)
+            self.labelIds = sorted(self.label2ind.keys())
+            self.num_cats = len(self.labelIds)
+
+            self.labelIds_base = buildLabelIndex(data_base['labels']).keys()
+            self.labelIds_novel = buildLabelIndex(data_novel['labels']).keys()
+            self.num_cats_base = len(self.labelIds_base)
+            self.num_cats_novel = len(self.labelIds_novel)
+            intersection = set(self.labelIds_base) & set(self.labelIds_novel)
+            assert(len(intersection) == 0)
+        else:
+            raise ValueError('Not valid phase {0}'.format(self.phase))
+
+        mean_pix = [x/255.0 for x in [129.37731888,
+                                      124.10583864, 112.47758569]]
+
+        std_pix = [x/255.0 for x in [68.20947949, 65.43124043, 70.45866994]]
+
+        normalize = transforms.Normalize(mean=mean_pix, std=std_pix)
+
+        if (self.phase == 'test' or self.phase == 'val') or (do_not_use_random_transf == True):
+
+            self.transform = transforms.Compose([
+                transforms.ToPILImage(),
+                # lambda x: np.asarray(x),
+                transforms.ToTensor(),
+                normalize
+            ])
+        else:
+            self.transform = transforms.Compose([
+                transforms.ToPILImage(),
+                transforms.RandomCrop(32, padding=4),
+                transforms.ColorJitter(
+                    brightness=0.4, contrast=0.4, saturation=0.4),
+                transforms.RandomHorizontalFlip(),
+                transforms.ToTensor(),
+                # lambda x: np.asarray(x),
+                normalize
+            ])
+
+    def __getitem__(self, index):
+        img, label = self.data[index], self.labels[index]
+        # doing this so that it is consistent with all other datasets
+        # to return a PIL Image
+
+        # img = Image.fromarray(img)
+        if self.transform is not None:
+            img = self.transform(img)
+        return img, label
+
+    def __len__(self):
+        return len(self.data)
+
+
+class FewShotDataloader():
+    def __init__(self,
+                 dataset,
+                 nKnovel=5,  # number of novel categories.
+                 nKbase=-1,  # number of base categories.
+                 # number of training examples per novel category.
+                 nExemplars=1,
+                 # number of test examples for all the novel categories.
+                 nTestNovel=15*5,
+                 # number of test examples for all the base categories.
+                 nTestBase=15*5,
+                 batch_size=1,  # number of training episodes per batch.
+                 num_workers=4,
+                 epoch_size=2000,  # number of batches per epoch.
+                 ):
+
+        self.dataset = dataset
+        self.phase = self.dataset.phase
+        max_possible_nKnovel = (self.dataset.num_cats_base if self.phase == 'train' or self.phase == 'trainval'
+                                else self.dataset.num_cats_novel)
+        assert(nKnovel >= 0 and nKnovel < max_possible_nKnovel)
+        self.nKnovel = nKnovel
+
+        max_possible_nKbase = self.dataset.num_cats_base
+        nKbase = nKbase if nKbase >= 0 else max_possible_nKbase
+        if (self.phase == 'train' or self.phase == 'trainval') and nKbase > 0:
+            nKbase -= self.nKnovel
+            max_possible_nKbase -= self.nKnovel
+
+        assert(nKbase >= 0 and nKbase <= max_possible_nKbase)
+        self.nKbase = nKbase
+
+        self.nExemplars = nExemplars
+        self.nTestNovel = nTestNovel
+        self.nTestBase = nTestBase
+        self.batch_size = batch_size
+        self.epoch_size = epoch_size
+        self.num_workers = num_workers
+        self.is_eval_mode = (self.phase == 'test') or (self.phase == 'val')
+
+    def sampleImageIdsFrom(self, cat_id, sample_size=1):
+        """
+        Samples `sample_size` number of unique image ids picked from the
+        category `cat_id` (i.e., self.dataset.label2ind[cat_id]).
+
+        Args:
+            cat_id: a scalar with the id of the category from which images will
+                be sampled.
+            sample_size: number of images that will be sampled.
+
+        Returns:
+            image_ids: a list of length `sample_size` with unique image ids.
+        """
+        assert(cat_id in self.dataset.label2ind)
+        assert(len(self.dataset.label2ind[cat_id]) >= sample_size)
+        # Note: random.sample samples elements without replacement.
+        # seed = random.randint(1,10000000)
+        # random.seed(seed)
+        return random.sample(self.dataset.label2ind[cat_id], sample_size)
+
+    def sampleCategories(self, cat_set, sample_size=1):
+        """
+        Samples `sample_size` number of unique categories picked from the
+        `cat_set` set of categories. `cat_set` can be either 'base' or 'novel'.
+
+        Args:
+            cat_set: string that specifies the set of categories from which
+                categories will be sampled.
+            sample_size: number of categories that will be sampled.
+
+        Returns:
+            cat_ids: a list of length `sample_size` with unique category ids.
+        """
+        if cat_set == 'base':
+            labelIds = self.dataset.labelIds_base
+        elif cat_set == 'novel':
+            labelIds = self.dataset.labelIds_novel
+        else:
+            raise ValueError('Not recognized category set {}'.format(cat_set))
+
+        assert(len(labelIds) >= sample_size)
+        # return sample_size unique categories chosen from labelIds set of
+        # categories (that can be either self.labelIds_base or self.labelIds_novel)
+        # Note: random.sample samples elements without replacement.
+        return random.sample(labelIds, sample_size)
+
+    def sample_base_and_novel_categories(self, nKbase, nKnovel):
+        """
+        Samples `nKbase` number of base categories and `nKnovel` number of novel
+        categories.
+
+        Args:
+            nKbase: number of base categories
+            nKnovel: number of novel categories
+
+        Returns:
+            Kbase: a list of length 'nKbase' with the ids of the sampled base
+                categories.
+            Knovel: a list of lenght 'nKnovel' with the ids of the sampled novel
+                categories.
+        """
+        if self.is_eval_mode:
+            assert(nKnovel <= self.dataset.num_cats_novel)
+            # sample from the set of base categories 'nKbase' number of base
+            # categories.
+            Kbase = sorted(self.sampleCategories('base', nKbase))
+            # sample from the set of novel categories 'nKnovel' number of novel
+            # categories.
+            Knovel = sorted(self.sampleCategories('novel', nKnovel))
+        else:
+            # sample from the set of base categories 'nKnovel' + 'nKbase' number
+            # of categories.
+            cats_ids = self.sampleCategories('base', nKnovel+nKbase)
+            assert(len(cats_ids) == (nKnovel+nKbase))
+            # Randomly pick 'nKnovel' number of fake novel categories and keep
+            # the rest as base categories.
+            random.shuffle(cats_ids)
+            Knovel = sorted(cats_ids[:nKnovel])
+            Kbase = sorted(cats_ids[nKnovel:])
+
+        return Kbase, Knovel
+
+    def sample_test_examples_for_base_categories(self, Kbase, nTestBase):
+        """
+        Sample `nTestBase` number of images from the `Kbase` categories.
+
+        Args:
+            Kbase: a list of length `nKbase` with the ids of the categories from
+                where the images will be sampled.
+            nTestBase: the total number of images that will be sampled.
+
+        Returns:
+            Tbase: a list of length `nTestBase` with 2-element tuples. The 1st
+                element of each tuple is the image id that was sampled and the
+                2nd elemend is its category label (which is in the range
+                [0, len(Kbase)-1]).
+        """
+        Tbase = []
+        if len(Kbase) > 0:
+            # Sample for each base category a number images such that the total
+            # number sampled images of all categories to be equal to `nTestBase`.
+            KbaseIndices = np.random.choice(
+                np.arange(len(Kbase)), size=nTestBase, replace=True)
+            KbaseIndices, NumImagesPerCategory = np.unique(
+                KbaseIndices, return_counts=True)
+
+            for Kbase_idx, NumImages in zip(KbaseIndices, NumImagesPerCategory):
+                imd_ids = self.sampleImageIdsFrom(
+                    Kbase[Kbase_idx], sample_size=NumImages)
+                Tbase += [(img_id, Kbase_idx) for img_id in imd_ids]
+
+        assert(len(Tbase) == nTestBase)
+
+        return Tbase
+
+    def sample_train_and_test_examples_for_novel_categories(
+            self, Knovel, nTestNovel, nExemplars, nKbase):
+        """Samples train and test examples of the novel categories.
+
+        Args:
+            Knovel: a list with the ids of the novel categories.
+            nTestNovel: the total number of test images that will be sampled
+                from all the novel categories.
+            nExemplars: the number of training examples per novel category that
+                will be sampled.
+            nKbase: the number of base categories. It is used as offset of the
+                category index of each sampled image.
+
+        Returns:
+            Tnovel: a list of length `nTestNovel` with 2-element tuples. The
+                1st element of each tuple is the image id that was sampled and
+                the 2nd element is its category label (which is in the range
+                [nKbase, nKbase + len(Knovel) - 1]).
+            Exemplars: a list of length len(Knovel) * nExemplars of 2-element
+                tuples. The 1st element of each tuple is the image id that was
+                sampled and the 2nd element is its category label (which is in
+                the ragne [nKbase, nKbase + len(Knovel) - 1]).
+        """
+
+        if len(Knovel) == 0:
+            return [], []
+
+        nKnovel = len(Knovel)
+        Tnovel = []
+        Exemplars = []
+        assert((nTestNovel % nKnovel) == 0)
+        nEvalExamplesPerClass = int(nTestNovel / nKnovel)
+
+        for Knovel_idx in range(nKnovel):
+            imd_ids = self.sampleImageIdsFrom(
+                Knovel[Knovel_idx],
+                sample_size=(nEvalExamplesPerClass + nExemplars))
+
+            imds_tnovel = imd_ids[:nEvalExamplesPerClass]
+            imds_ememplars = imd_ids[nEvalExamplesPerClass:]
+
+            Tnovel += [(img_id, nKbase+Knovel_idx) for img_id in imds_tnovel]
+
+            Exemplars += [(img_id, nKbase+Knovel_idx)
+                          for img_id in imds_ememplars]
+
+        # print('='*60)
+        # print(Tnovel)
+        # print(Exemplars)
+        # print('='*60)
+        assert(len(Tnovel) == nTestNovel)
+        assert(len(Exemplars) == len(Knovel) * nExemplars)
+
+        # random.shuffle(Exemplars)  # shuffle commented by me
+
+        # print(Exemplars)
+
+        return Tnovel, Exemplars
+
+    def sample_episode(self):
+        """Samples a training episode."""
+        nKnovel = self.nKnovel
+        nKbase = self.nKbase
+        nTestNovel = self.nTestNovel
+        nTestBase = self.nTestBase
+        nExemplars = self.nExemplars
+
+        Kbase, Knovel = self.sample_base_and_novel_categories(nKbase, nKnovel)
+
+        Tbase = self.sample_test_examples_for_base_categories(Kbase, nTestBase)
+
+        # print(Tbase,Knovel)
+
+        Tnovel, Exemplars = self.sample_train_and_test_examples_for_novel_categories(
+            Knovel, nTestNovel, nExemplars, nKbase)
+
+        # concatenate the base and novel category examples.
+        Test = Tbase + Tnovel
+        # random.shuffle(Test)
+
+        # print(Test)
+
+        Kall = Kbase + Knovel
+
+        return Exemplars, Test, Kall, nKbase
+
+    def createExamplesTensorData(self, examples):
+        """
+        Creates the examples image and label tensor data.
+
+        Args:
+            examples: a list of 2-element tuples, each representing a
+                train or test example. The 1st element of each tuple
+                is the image id of the example and 2nd element is the
+                category label of the example, which is in the range
+                [0, nK - 1], where nK is the total number of categories
+                (both novel and base).
+
+        Returns:
+            images: a tensor of shape [nExamples, Height, Width, 3] with the
+                example images, where nExamples is the number of examples
+                (i.e., nExamples = len(examples)).
+            labels: a tensor of shape [nExamples] with the category label
+                of each example.
+        """
+        images = torch.stack(
+            [self.dataset[img_idx][0] for img_idx, _ in examples], dim=0)
+        labels = torch.LongTensor([label for _, label in examples])
+        return images, labels
+
+    def get_iterator(self, epoch=0):
+        rand_seed = epoch
+        random.seed(rand_seed)
+        np.random.seed(rand_seed)
+
+        def load_function(iter_idx):
+            Exemplars, Test, Kall, nKbase = self.sample_episode()
+            Xt, Yt = self.createExamplesTensorData(Test)
+            Kall = torch.LongTensor(Kall)
+            if len(Exemplars) > 0:
+                Xe, Ye = self.createExamplesTensorData(Exemplars)
+                return Xe, Ye, Xt, Yt, Kall, nKbase
+            else:
+                return Xt, Yt, Kall, nKbase
+
+        tnt_dataset = tnt.dataset.ListDataset(
+            elem_list=range(self.epoch_size), load=load_function)
+        data_loader = tnt_dataset.parallel(
+            batch_size=self.batch_size,
+            num_workers=(1 if self.is_eval_mode else self.num_workers),
+            shuffle=(False if self.is_eval_mode else True))
+
+        return data_loader
+
+    def __call__(self, epoch=0):
+        return self.get_iterator(epoch)
+
+    def __len__(self):
+        return int(self.epoch_size / self.batch_size)

dataloader/FC100.py

@@ -0,0 +1,453 @@
+# Dataloader of Gidaris & Komodakis, CVPR 2018
+# Adapted from:
+# https://github.com/gidariss/FewShotWithoutForgetting/blob/master/dataloader.py
+from __future__ import print_function
+
+import os
+import os.path
+import numpy as np
+import random
+import pickle
+import json
+import math
+
+import torch
+import torch.utils.data as data
+import torchvision
+import torchvision.datasets as datasets
+import torchvision.transforms as transforms
+import torchnet as tnt
+
+import h5py
+
+from PIL import Image
+from PIL import ImageEnhance
+
+from pdb import set_trace as breakpoint
+
+
+# Set the appropriate paths of the datasets here.
+_FC100_DATASET_DIR = './cifar/FC100/'
+
+def buildLabelIndex(labels):
+    label2inds = {}
+    for idx, label in enumerate(labels):
+        if label not in label2inds:
+            label2inds[label] = []
+        label2inds[label].append(idx)
+
+    return label2inds
+
+def load_data(file):
+    try:
+        with open(file, 'rb') as fo:
+            data = pickle.load(fo)
+        return data
+    except:
+        with open(file, 'rb') as f:
+            u = pickle._Unpickler(f)
+            u.encoding = 'latin1'
+            data = u.load()
+        return data
+
+class FC100(data.Dataset):
+    def __init__(self, phase='train', do_not_use_random_transf=False):
+
+        assert(phase=='train' or phase=='val' or phase=='test'or phase=='trainval')
+        self.phase = phase
+        self.name = 'FC100_' + phase
+
+        print('Loading FC100 dataset - phase {0}'.format(phase))
+        file_train_categories_train_phase = os.path.join(
+            _FC100_DATASET_DIR,
+            'FC100_train.pickle')
+        file_train_categories_val_phase = os.path.join(
+            _FC100_DATASET_DIR,
+            'FC100_train.pickle')
+        file_train_categories_test_phase = os.path.join(
+            _FC100_DATASET_DIR,
+            'FC100_train.pickle')
+        file_val_categories_val_phase = os.path.join(
+            _FC100_DATASET_DIR,
+            'FC100_val.pickle')
+        file_test_categories_test_phase = os.path.join(
+            _FC100_DATASET_DIR,
+            'FC100_test.pickle')
+
+        if self.phase=='train':
+            # During training phase we only load the training phase images
+            # of the training categories (aka base categories).
+            data_train = load_data(file_train_categories_train_phase)
+            self.data = data_train['data']
+            self.labels = data_train['labels']
+
+            #print (self.labels)
+            self.label2ind = buildLabelIndex(self.labels)
+            self.labelIds = sorted(self.label2ind.keys())
+            self.num_cats = len(self.labelIds)
+            self.labelIds_base = self.labelIds
+            self.num_cats_base = len(self.labelIds_base)
+            #print (self.data.shape)
+        elif self.phase == 'trainval':
+            # During training phase we only load the training phase images
+            # of the training categories (aka base categories).
+            data_train = load_data(file_train_categories_train_phase)
+            self.data = data_train['data']
+            self.labels = data_train['labels']
+            data_base = load_data(file_train_categories_val_phase)
+            data_novel = load_data(file_val_categories_val_phase)
+            self.data = np.concatenate(
+                [self.data, data_novel['data']], axis=0)
+            self.data = np.concatenate(
+                [self.data, data_base['data']], axis=0)
+
+            self.labels = np.concatenate(
+                [self.labels, data_novel['labels']], axis=0)
+            self.labels = np.concatenate(
+                [self.labels, data_base['labels']], axis=0)
+
+            # print (self.labels)
+            self.label2ind = buildLabelIndex(self.labels)
+            self.labelIds = sorted(self.label2ind.keys())
+            self.num_cats = len(self.labelIds)
+            self.labelIds_base = self.labelIds
+            self.num_cats_base = len(self.labelIds_base)
+        elif self.phase=='val' or self.phase=='test':
+            if self.phase=='test':
+                # load data that will be used for evaluating the recognition
+                # accuracy of the base categories.
+                data_base = load_data(file_train_categories_test_phase)
+                # load data that will be use for evaluating the few-shot recogniton
+                # accuracy on the novel categories.
+                data_novel = load_data(file_test_categories_test_phase)
+            else: # phase=='val'
+                # load data that will be used for evaluating the recognition
+                # accuracy of the base categories.
+                data_base = load_data(file_train_categories_val_phase)
+                # load data that will be use for evaluating the few-shot recogniton
+                # accuracy on the novel categories.
+                data_novel = load_data(file_val_categories_val_phase)
+
+            self.data = np.concatenate(
+                [data_base['data'], data_novel['data']], axis=0)
+            self.labels = data_base['labels'] + data_novel['labels']
+
+            self.label2ind = buildLabelIndex(self.labels)
+            self.labelIds = sorted(self.label2ind.keys())
+            self.num_cats = len(self.labelIds)
+
+            self.labelIds_base = buildLabelIndex(data_base['labels']).keys()
+            self.labelIds_novel = buildLabelIndex(data_novel['labels']).keys()
+            self.num_cats_base = len(self.labelIds_base)
+            self.num_cats_novel = len(self.labelIds_novel)
+            intersection = set(self.labelIds_base) & set(self.labelIds_novel)
+            assert(len(intersection) == 0)
+        else:
+            raise ValueError('Not valid phase {0}'.format(self.phase))
+
+        mean_pix = [x/255.0 for x in [129.37731888, 124.10583864, 112.47758569]]
+
+        std_pix = [x/255.0 for x in [68.20947949, 65.43124043, 70.45866994]]
+        
+        normalize = transforms.Normalize(mean=mean_pix, std=std_pix)
+
+        if (self.phase=='test' or self.phase=='val') or (do_not_use_random_transf==True):
+            self.transform = transforms.Compose([
+                lambda x: np.asarray(x),
+                transforms.ToTensor(),
+                normalize
+            ])
+        else:
+            self.transform = transforms.Compose([
+                transforms.RandomCrop(32, padding=4),
+                transforms.ColorJitter(brightness=0.4, contrast=0.4, saturation=0.4),
+                transforms.RandomHorizontalFlip(),
+                lambda x: np.asarray(x),
+                transforms.ToTensor(),
+                normalize
+            ])
+
+    def __getitem__(self, index):
+        img, label = self.data[index], self.labels[index]
+        # doing this so that it is consistent with all other datasets
+        # to return a PIL Image
+        img = Image.fromarray(img)
+        if self.transform is not None:
+            img = self.transform(img)
+        return img, label
+
+    def __len__(self):
+        return len(self.data)
+
+
+class FewShotDataloader():
+    def __init__(self,
+                 dataset,
+                 nKnovel=5, # number of novel categories.
+                 nKbase=-1, # number of base categories.
+                 nExemplars=1, # number of training examples per novel category.
+                 nTestNovel=15*5, # number of test examples for all the novel categories.
+                 nTestBase=15*5, # number of test examples for all the base categories.
+                 batch_size=1, # number of training episodes per batch.
+                 num_workers=4,
+                 epoch_size=2000, # number of batches per epoch.
+                 ):
+
+        self.dataset = dataset
+        self.phase = self.dataset.phase
+        max_possible_nKnovel = (self.dataset.num_cats_base if self.phase=='train' or self.phase=='trainval'
+                                else self.dataset.num_cats_novel)
+        assert(nKnovel >= 0 and nKnovel < max_possible_nKnovel)
+        self.nKnovel = nKnovel
+
+        max_possible_nKbase = self.dataset.num_cats_base
+        nKbase = nKbase if nKbase >= 0 else max_possible_nKbase
+        if (self.phase=='train' or self.phase=='trainval') and nKbase > 0:
+            nKbase -= self.nKnovel
+            max_possible_nKbase -= self.nKnovel
+
+        assert(nKbase >= 0 and nKbase <= max_possible_nKbase)
+        self.nKbase = nKbase
+
+        self.nExemplars = nExemplars
+        self.nTestNovel = nTestNovel
+        self.nTestBase = nTestBase
+        self.batch_size = batch_size
+        self.epoch_size = epoch_size
+        self.num_workers = num_workers
+        self.is_eval_mode = (self.phase=='test') or (self.phase=='val')
+
+    def sampleImageIdsFrom(self, cat_id, sample_size=1):
+        """
+        Samples `sample_size` number of unique image ids picked from the
+        category `cat_id` (i.e., self.dataset.label2ind[cat_id]).
+
+        Args:
+            cat_id: a scalar with the id of the category from which images will
+                be sampled.
+            sample_size: number of images that will be sampled.
+
+        Returns:
+            image_ids: a list of length `sample_size` with unique image ids.
+        """
+        assert(cat_id in self.dataset.label2ind)
+        assert(len(self.dataset.label2ind[cat_id]) >= sample_size)
+        # Note: random.sample samples elements without replacement.
+        return random.sample(self.dataset.label2ind[cat_id], sample_size)
+
+    def sampleCategories(self, cat_set, sample_size=1):
+        """
+        Samples `sample_size` number of unique categories picked from the
+        `cat_set` set of categories. `cat_set` can be either 'base' or 'novel'.
+
+        Args:
+            cat_set: string that specifies the set of categories from which
+                categories will be sampled.
+            sample_size: number of categories that will be sampled.
+
+        Returns:
+            cat_ids: a list of length `sample_size` with unique category ids.
+        """
+        if cat_set=='base':
+            labelIds = self.dataset.labelIds_base
+        elif cat_set=='novel':
+            labelIds = self.dataset.labelIds_novel
+        else:
+            raise ValueError('Not recognized category set {}'.format(cat_set))
+
+        assert(len(labelIds) >= sample_size)
+        # return sample_size unique categories chosen from labelIds set of
+        # categories (that can be either self.labelIds_base or self.labelIds_novel)
+        # Note: random.sample samples elements without replacement.
+        return random.sample(labelIds, sample_size)
+
+    def sample_base_and_novel_categories(self, nKbase, nKnovel):
+        """
+        Samples `nKbase` number of base categories and `nKnovel` number of novel
+        categories.
+
+        Args:
+            nKbase: number of base categories
+            nKnovel: number of novel categories
+
+        Returns:
+            Kbase: a list of length 'nKbase' with the ids of the sampled base
+                categories.
+            Knovel: a list of lenght 'nKnovel' with the ids of the sampled novel
+                categories.
+        """
+        if self.is_eval_mode:
+            assert(nKnovel <= self.dataset.num_cats_novel)
+            # sample from the set of base categories 'nKbase' number of base
+            # categories.
+            Kbase = sorted(self.sampleCategories('base', nKbase))
+            # sample from the set of novel categories 'nKnovel' number of novel
+            # categories.
+            Knovel = sorted(self.sampleCategories('novel', nKnovel))
+        else:
+            # sample from the set of base categories 'nKnovel' + 'nKbase' number
+            # of categories.
+            cats_ids = self.sampleCategories('base', nKnovel+nKbase)
+            assert(len(cats_ids) == (nKnovel+nKbase))
+            # Randomly pick 'nKnovel' number of fake novel categories and keep
+            # the rest as base categories.
+            random.shuffle(cats_ids)
+            Knovel = sorted(cats_ids[:nKnovel])
+            Kbase = sorted(cats_ids[nKnovel:])
+
+        return Kbase, Knovel
+
+    def sample_test_examples_for_base_categories(self, Kbase, nTestBase):
+        """
+        Sample `nTestBase` number of images from the `Kbase` categories.
+
+        Args:
+            Kbase: a list of length `nKbase` with the ids of the categories from
+                where the images will be sampled.
+            nTestBase: the total number of images that will be sampled.
+
+        Returns:
+            Tbase: a list of length `nTestBase` with 2-element tuples. The 1st
+                element of each tuple is the image id that was sampled and the
+                2nd elemend is its category label (which is in the range
+                [0, len(Kbase)-1]).
+        """
+        Tbase = []
+        if len(Kbase) > 0:
+            # Sample for each base category a number images such that the total
+            # number sampled images of all categories to be equal to `nTestBase`.
+            KbaseIndices = np.random.choice(
+                np.arange(len(Kbase)), size=nTestBase, replace=True)
+            KbaseIndices, NumImagesPerCategory = np.unique(
+                KbaseIndices, return_counts=True)
+
+            for Kbase_idx, NumImages in zip(KbaseIndices, NumImagesPerCategory):
+                imd_ids = self.sampleImageIdsFrom(
+                    Kbase[Kbase_idx], sample_size=NumImages)
+                Tbase += [(img_id, Kbase_idx) for img_id in imd_ids]
+
+        assert(len(Tbase) == nTestBase)
+
+        return Tbase
+
+    def sample_train_and_test_examples_for_novel_categories(
+            self, Knovel, nTestNovel, nExemplars, nKbase):
+        """Samples train and test examples of the novel categories.
+
+        Args:
+    	    Knovel: a list with the ids of the novel categories.
+            nTestNovel: the total number of test images that will be sampled
+                from all the novel categories.
+            nExemplars: the number of training examples per novel category that
+                will be sampled.
+            nKbase: the number of base categories. It is used as offset of the
+                category index of each sampled image.
+
+        Returns:
+            Tnovel: a list of length `nTestNovel` with 2-element tuples. The
+                1st element of each tuple is the image id that was sampled and
+                the 2nd element is its category label (which is in the range
+                [nKbase, nKbase + len(Knovel) - 1]).
+            Exemplars: a list of length len(Knovel) * nExemplars of 2-element
+                tuples. The 1st element of each tuple is the image id that was
+                sampled and the 2nd element is its category label (which is in
+                the ragne [nKbase, nKbase + len(Knovel) - 1]).
+        """
+
+        if len(Knovel) == 0:
+            return [], []
+
+        nKnovel = len(Knovel)
+        Tnovel = []
+        Exemplars = []
+        assert((nTestNovel % nKnovel) == 0)
+        nEvalExamplesPerClass = int(nTestNovel / nKnovel)
+
+        for Knovel_idx in range(len(Knovel)):
+            imd_ids = self.sampleImageIdsFrom(
+                Knovel[Knovel_idx],
+                sample_size=(nEvalExamplesPerClass + nExemplars))
+
+            imds_tnovel = imd_ids[:nEvalExamplesPerClass]
+            imds_ememplars = imd_ids[nEvalExamplesPerClass:]
+
+            Tnovel += [(img_id, nKbase+Knovel_idx) for img_id in imds_tnovel]
+            Exemplars += [(img_id, nKbase+Knovel_idx) for img_id in imds_ememplars]
+        assert(len(Tnovel) == nTestNovel)
+        assert(len(Exemplars) == len(Knovel) * nExemplars)
+        random.shuffle(Exemplars)
+
+        return Tnovel, Exemplars
+
+    def sample_episode(self):
+        """Samples a training episode."""
+        nKnovel = self.nKnovel
+        nKbase = self.nKbase
+        nTestNovel = self.nTestNovel
+        nTestBase = self.nTestBase
+        nExemplars = self.nExemplars
+
+        Kbase, Knovel = self.sample_base_and_novel_categories(nKbase, nKnovel)
+        Tbase = self.sample_test_examples_for_base_categories(Kbase, nTestBase)
+        Tnovel, Exemplars = self.sample_train_and_test_examples_for_novel_categories(
+            Knovel, nTestNovel, nExemplars, nKbase)
+
+        # concatenate the base and novel category examples.
+        Test = Tbase + Tnovel
+        random.shuffle(Test)
+        Kall = Kbase + Knovel
+
+        return Exemplars, Test, Kall, nKbase
+
+    def createExamplesTensorData(self, examples):
+        """
+        Creates the examples image and label tensor data.
+
+        Args:
+            examples: a list of 2-element tuples, each representing a
+                train or test example. The 1st element of each tuple
+                is the image id of the example and 2nd element is the
+                category label of the example, which is in the range
+                [0, nK - 1], where nK is the total number of categories
+                (both novel and base).
+
+        Returns:
+            images: a tensor of shape [nExamples, Height, Width, 3] with the
+                example images, where nExamples is the number of examples
+                (i.e., nExamples = len(examples)).
+            labels: a tensor of shape [nExamples] with the category label
+                of each example.
+        """
+        images = torch.stack(
+            [self.dataset[img_idx][0] for img_idx, _ in examples], dim=0)
+        labels = torch.LongTensor([label for _, label in examples])
+        return images, labels
+
+    def get_iterator(self, epoch=0):
+        rand_seed = epoch
+        random.seed(rand_seed)
+        np.random.seed(rand_seed)
+        def load_function(iter_idx):
+            Exemplars, Test, Kall, nKbase = self.sample_episode()
+            Xt, Yt = self.createExamplesTensorData(Test)
+            Kall = torch.LongTensor(Kall)
+            if len(Exemplars) > 0:
+                Xe, Ye = self.createExamplesTensorData(Exemplars)
+                return Xe, Ye, Xt, Yt, Kall, nKbase
+            else:
+                return Xt, Yt, Kall, nKbase
+
+        tnt_dataset = tnt.dataset.ListDataset(
+            elem_list=range(self.epoch_size), load=load_function)
+        data_loader = tnt_dataset.parallel(
+            batch_size=self.batch_size,
+            num_workers=(0 if self.is_eval_mode else self.num_workers),
+            shuffle=(False if self.is_eval_mode else True))
+
+        return data_loader
+
+    def __call__(self, epoch=0):
+        return self.get_iterator(epoch)
+
+    def __len__(self):
+        return int(self.epoch_size / self.batch_size)

dataloader/chest.py

@@ -0,0 +1,512 @@
+# Dataloader of Gidaris & Komodakis, CVPR 2018
+# Adapted from:
+# https://github.com/gidariss/FewShotWithoutForgetting/blob/master/dataloader.py
+from __future__ import print_function
+
+import os
+import os.path
+import numpy as npw
+import random
+import pickle
+import json
+import math
+
+import torch
+import torch.utils.data as data
+import torchvision
+import torchvision.datasets as datasets
+import torchvision.transforms as transforms
+import torchnet as tnt
+import numpy as np
+import pandas as pd
+
+
+import h5py
+
+import cv2
+from PIL import Image
+from PIL import ImageEnhance
+import matplotlib.pyplot as plt
+
+
+from torchvision.transforms.transforms import ToPILImage
+
+
+# Set the appropriate paths of the datasets here.
+# _CIFAR_FS_DATASET_DIR = './cifar/CIFAR-FS/'
+_CHEST_DATASET_DIR = './NIH'
+image_path = './NIH/images'
+
+
+label_dict = {'Cardiomegaly': 0, 'Edema': 1, 'Effusion': 2, 'Emphysema': 3, 'Infiltration': 4, 'Mass': 5, 'Atelectasis': 6, 'Consolidation': 7,
+              'Pleural_Thickening': 8, 'Fibrosis': 9, 'Hernia': 10, 'Pneumonia': 11, 'Nodule': 12, 'Pneumothorax': 13, 'No Finding': 14}
+
+
+def buildLabelIndex(labels):
+    label2inds = {}
+    for idx, label in enumerate(labels):
+        label = label_dict[label]
+        if label not in label2inds:
+            label2inds[label] = []
+        label2inds[label].append(idx)
+
+    return label2inds
+
+
+def load_data(file):
+    try:
+        with open(file, 'rb') as fo:
+            data = pickle.load(fo)
+        return data
+    except:
+        with open(file, 'rb') as f:
+            u = pickle._Unpickler(f)
+            u.encoding = 'latin1'
+            data = u.load()
+        return data
+
+
+class Chest(data.Dataset):
+    def __init__(self, phase='train', idx = 1, do_not_use_random_transf=False):
+
+        assert(phase == 'train' or phase == 'val' or phase ==
+               'test' or phase == 'trainval')
+        self.phase = phase
+        # self.name = phase + '.csv'
+
+        # idx = 3 # represents group for experimentation
+
+        print('Loading Chest-XRay dataset - phase {0}'.format(phase))
+
+        train_path = os.path.join(_CHEST_DATASET_DIR, f'train{idx}.csv')
+        val_path = os.path.join(_CHEST_DATASET_DIR, f'val{idx}.csv')
+        test_path = os.path.join(_CHEST_DATASET_DIR, f'test{idx}.csv')
+
+        if self.phase == 'train':
+            # # During training phase we only load the training phase images
+            # # of the training categories (aka base categories).
+            # data_train = load_data(file_train_categories_train_phase)
+            # # self.data = data_train['data']
+            # self.labels = data_train['labels']
+
+            file = pd.read_csv(train_path)
+
+            self.data = file['image_id'].values
+
+            self.labels = file['class_name'].values
+
+            self.label2ind = buildLabelIndex(self.labels)
+
+            self.labelIds = sorted(self.label2ind.keys())
+            self.num_cats = len(self.labelIds)
+            self.labelIds_base = self.labelIds
+            self.num_cats_base = len(self.labelIds_base)
+
+        # elif self.phase == 'trainval':
+        #     # During training phase we only load the training phase images
+        #     # of the training categories (aka base categories).
+        #     data_train = load_data(file_train_categories_train_phase)
+        #     self.data = data_train['data']
+        #     self.labels = data_train['labels']
+        #     data_base = load_data(file_train_categories_val_phase)
+        #     data_novel = load_data(file_val_categories_val_phase)
+        #     self.data = np.concatenate(
+        #         [self.data, data_novel['data']], axis=0)
+        #     self.data = np.concatenate(
+        #         [self.data, data_base['data']], axis=0)
+
+        #     self.labels = np.concatenate(
+        #         [self.labels, data_novel['labels']], axis=0)
+        #     self.labels = np.concatenate(
+        #         [self.labels, data_base['labels']], axis=0)
+
+        #     self.label2ind = buildLabelIndex(self.labels)
+        #     self.labelIds = sorted(self.label2ind.keys())
+        #     self.num_cats = len(self.labelIds)
+        #     self.labelIds_base = self.labelIds
+        #     self.num_cats_base = len(self.labelIds_base)
+
+        elif self.phase == 'val' or self.phase == 'test':
+            if self.phase == 'test':
+                # # load data that will be used for evaluating the recognition
+                # # accuracy of the base categories.
+                # data_base = load_data(file_train_categories_test_phase)
+                # # load data that will be use for evaluating the few-shot recogniton
+                # # accuracy on the novel categories.
+                # data_novel = load_data(file_test_categories_test_phase)
+
+                train_file = pd.read_csv(train_path)
+                file = pd.read_csv(test_path)
+            else:  # phase=='val'
+                # # load data that will be used for evaluating the recognition
+                # # accuracy of the base categories.
+                # data_base = load_data(file_train_categories_val_phase)
+                # # load data that will be use for evaluating the few-shot recogniton
+                # # accuracy on the novel categories.
+                # data_novel = load_data(file_val_categories_val_phase)
+
+                train_file = pd.read_csv(train_path)
+                file = pd.read_csv(val_path)
+
+            # self.data = np.concatenate(
+            #     [data_base['data'], data_novel['data']], axis=0)
+            # self.labels = data_base['labels'] + data_novel['labels']
+
+            train_labels = train_file['class_name'].values
+            novel_labels = file['class_name'].values
+
+            self.data = np.concatenate(
+                [train_file['image_id'].values, file['image_id'].values], axis=0)
+            self.labels = np.concatenate(
+                [train_file['class_name'].values, file['class_name'].values], axis=0)
+
+            self.label2ind = buildLabelIndex(self.labels)
+            self.labelIds = sorted(self.label2ind.keys())
+            self.num_cats = len(self.labelIds)
+
+            # self.labelIds_base = buildLabelIndex(data_base['labels']).keys()
+            # self.labelIds_novel = buildLabelIndex(data_novel['labels']).keys()
+
+            self.labelIds_base = buildLabelIndex(train_labels).keys()
+            self.labelIds_novel = buildLabelIndex(novel_labels).keys()
+            print('='*60)
+            print(self.labelIds_novel)
+            print('='*60)
+
+            self.num_cats_base = len(self.labelIds_base)
+            self.num_cats_novel = len(self.labelIds_novel)
+            # print(self.labelIds_novel)
+            # print(self.num_cats_novel)
+            intersection = set(self.labelIds_base) & set(self.labelIds_novel)
+            assert(len(intersection) == 0)
+        else:
+            raise ValueError('Not valid phase {0}'.format(self.phase))
+
+        # mean_pix = [x/255.0 for x in [129.37731888,
+        #                               124.10583864, 112.47758569]]
+
+        # std_pix = [x/255.0 for x in [68.20947949, 65.43124043, 70.45866994]]
+
+        mean_pix =  [0.52024849, 0.52024849, 0.52024849]
+        std_pix =  [0.22699496, 0.22699496, 0.22699496]
+
+
+        normalize = transforms.Normalize(mean=mean_pix, std=std_pix)
+
+        if (self.phase == 'test' or self.phase == 'val') or (do_not_use_random_transf == True):
+
+            self.transform = transforms.Compose([
+                transforms.ToPILImage(),
+                # lambda x: np.asarray(x),
+                transforms.ToTensor(),
+                # lambda x: x/255.0,
+                normalize
+            ])
+        else:
+            self.transform = transforms.Compose([
+                transforms.ToPILImage(),
+                # transforms.RandomCrop(32, padding=4),
+                # transforms.ColorJitter(
+                #     brightness=0.4, contrast=0.4, saturation=0.4),
+                transforms.RandomHorizontalFlip(),
+                transforms.ToTensor(),
+                # lambda x: np.asarray(x),
+                # lambda x: x/255.0,
+                normalize
+            ])
+
+    def __getitem__(self, index):
+        img, label = cv2.imread(os.path.join(
+            image_path, self.data[index]))[:,:,::-1], self.labels[index]
+        img = cv2.resize(img,(128,128)) # resize by Garvit
+        # img = cv2.resize(img,(84, 84)) # resize by kshitiz
+
+        # img = Image.fromarray(img)
+        if self.transform is not None:
+            img = self.transform(img)
+        return img, label
+
+    def __len__(self):
+        return len(self.data)
+
+
+class FewShotDataloader():
+    def __init__(self,
+                 dataset,
+                 nKnovel=5,  # number of novel categories.
+                 nKbase=-1,  # number of base categories.
+                 # number of training examples per novel category.
+                 nExemplars=1,
+                 # number of test examples for all the novel categories.
+                 nTestNovel=15*5,
+                 # number of test examples for all the base categories.
+                 nTestBase=15*5,
+                 batch_size=1,  # number of training episodes per batch.
+                 num_workers=4,
+                 epoch_size=2000,  # number of batches per epoch.
+                 ):
+
+        self.dataset = dataset
+        self.phase = self.dataset.phase
+        max_possible_nKnovel = (self.dataset.num_cats_base if self.phase == 'train' or self.phase == 'trainval'
+                                else self.dataset.num_cats_novel)
+
+        assert(nKnovel >= 0 and nKnovel <= max_possible_nKnovel)
+        self.nKnovel = nKnovel
+
+        max_possible_nKbase = self.dataset.num_cats_base
+        nKbase = nKbase if nKbase >= 0 else max_possible_nKbase
+        if (self.phase == 'train' or self.phase == 'trainval') and nKbase > 0:
+            nKbase -= self.nKnovel
+            max_possible_nKbase -= self.nKnovel
+
+        assert(nKbase >= 0 and nKbase <= max_possible_nKbase)
+        self.nKbase = nKbase
+
+        self.nExemplars = nExemplars
+        self.nTestNovel = nTestNovel
+        self.nTestBase = nTestBase
+        self.batch_size = batch_size
+        self.epoch_size = epoch_size
+        self.num_workers = num_workers
+        self.is_eval_mode = (self.phase == 'test') or (self.phase == 'val')
+
+    def sampleImageIdsFrom(self, cat_id, sample_size=1):
+        """
+        Samples `sample_size` number of unique image ids picked from the
+        category `cat_id` (i.e., self.dataset.label2ind[cat_id]).
+
+        Args:
+            cat_id: a scalar with the id of the category from which images will
+                be sampled.
+            sample_size: number of images that will be sampled.
+
+        Returns:
+            image_ids: a list of length `sample_size` with unique image ids.
+        """
+        assert(cat_id in self.dataset.label2ind)
+        assert(len(self.dataset.label2ind[cat_id]) >= sample_size)
+        # Note: random.sample samples elements without replacement.
+        # seed = random.randint(1,10000000)
+        # random.seed(seed)
+        return random.sample(self.dataset.label2ind[cat_id], sample_size)
+
+    def sampleCategories(self, cat_set, sample_size=1):
+        """
+        Samples `sample_size` number of unique categories picked from the
+        `cat_set` set of categories. `cat_set` can be either 'base' or 'novel'.
+
+        Args:
+            cat_set: string that specifies the set of categories from which
+                categories will be sampled.
+            sample_size: number of categories that will be sampled.
+
+        Returns:
+            cat_ids: a list of length `sample_size` with unique category ids.
+        """
+        if cat_set == 'base':
+            labelIds = self.dataset.labelIds_base
+        elif cat_set == 'novel':
+            labelIds = self.dataset.labelIds_novel
+        else:
+            raise ValueError('Not recognized category set {}'.format(cat_set))
+
+        assert(len(labelIds) >= sample_size)
+        # return sample_size unique categories chosen from labelIds set of
+        # categories (that can be either self.labelIds_base or self.labelIds_novel)
+        # Note: random.sample samples elements without replacement.
+        return random.sample(labelIds, sample_size)
+
+    def sample_base_and_novel_categories(self, nKbase, nKnovel):
+        """
+        Samples `nKbase` number of base categories and `nKnovel` number of novel
+        categories.
+
+        Args:
+            nKbase: number of base categories
+            nKnovel: number of novel categories
+
+        Returns:
+            Kbase: a list of length 'nKbase' with the ids of the sampled base
+                categories.
+            Knovel: a list of lenght 'nKnovel' with the ids of the sampled novel
+                categories.
+        """
+        if self.is_eval_mode:
+            assert(nKnovel <= self.dataset.num_cats_novel)
+            # sample from the set of base categories 'nKbase' number of base
+            # categories.
+            Kbase = sorted(self.sampleCategories('base', nKbase))
+            # sample from the set of novel categories 'nKnovel' number of novel
+            # categories.
+            Knovel = sorted(self.sampleCategories('novel', nKnovel))
+        else:
+            # sample from the set of base categories 'nKnovel' + 'nKbase' number
+            # of categories.
+            cats_ids = self.sampleCategories('base', nKnovel+nKbase)
+            assert(len(cats_ids) == (nKnovel+nKbase))
+            # Randomly pick 'nKnovel' number of fake novel categories and keep
+            # the rest as base categories.
+            random.shuffle(cats_ids)
+            Knovel = sorted(cats_ids[:nKnovel])
+            Kbase = sorted(cats_ids[nKnovel:])
+
+
+        return Kbase, Knovel
+
+    def sample_test_examples_for_base_categories(self, Kbase, nTestBase):
+        """
+        Sample `nTestBase` number of images from the `Kbase` categories.
+
+        Args:
+            Kbase: a list of length `nKbase` with the ids of the categories from
+                where the images will be sampled.
+            nTestBase: the total number of images that will be sampled.
+
+        Returns:
+            Tbase: a list of length `nTestBase` with 2-element tuples. The 1st
+                element of each tuple is the image id that was sampled and the
+                2nd elemend is its category label (which is in the range
+                [0, len(Kbase)-1]).
+        """
+        Tbase = []
+        if len(Kbase) > 0:
+            # Sample for each base category a number images such that the total
+            # number sampled images of all categories to be equal to `nTestBase`.
+            KbaseIndices = np.random.choice(
+                np.arange(len(Kbase)), size=nTestBase, replace=True)
+            KbaseIndices, NumImagesPerCategory = np.unique(
+                KbaseIndices, return_counts=True)
+
+            for Kbase_idx, NumImages in zip(KbaseIndices, NumImagesPerCategory):
+                imd_ids = self.sampleImageIdsFrom(
+                    Kbase[Kbase_idx], sample_size=NumImages)
+                Tbase += [(img_id, Kbase_idx) for img_id in imd_ids]
+
+        assert(len(Tbase) == nTestBase)
+
+        return Tbase
+
+    def sample_train_and_test_examples_for_novel_categories(
+            self, Knovel, nTestNovel, nExemplars, nKbase):
+        """Samples train and test examples of the novel categories.
+
+        Args:
+            Knovel: a list with the ids of the novel categories.
+            nTestNovel: the total number of test images that will be sampled
+                from all the novel categories.
+            nExemplars: the number of training examples per novel category that
+                will be sampled.
+            nKbase: the number of base categories. It is used as offset of the
+                category index of each sampled image.
+
+        Returns:
+            Tnovel: a list of length `nTestNovel` with 2-element tuples. The
+                1st element of each tuple is the image id that was sampled and
+                the 2nd element is its category label (which is in the range
+                [nKbase, nKbase + len(Knovel) - 1]).
+            Exemplars: a list of length len(Knovel) * nExemplars of 2-element
+                tuples. The 1st element of each tuple is the image id that was
+                sampled and the 2nd element is its category label (which is in
+                the ragne [nKbase, nKbase + len(Knovel) - 1]).
+        """
+
+        if len(Knovel) == 0:
+            return [], []
+
+        nKnovel = len(Knovel)
+        Tnovel = []
+        Exemplars = []
+        assert((nTestNovel % nKnovel) == 0)
+        nEvalExamplesPerClass = int(nTestNovel / nKnovel)
+
+        for Knovel_idx in range(len(Knovel)):
+            imd_ids = self.sampleImageIdsFrom(
+                Knovel[Knovel_idx],
+                sample_size=(nEvalExamplesPerClass + nExemplars))
+
+            imds_tnovel = imd_ids[:nEvalExamplesPerClass]
+            imds_ememplars = imd_ids[nEvalExamplesPerClass:]
+
+            Tnovel += [(img_id, nKbase+Knovel_idx) for img_id in imds_tnovel]
+            Exemplars += [(img_id, nKbase+Knovel_idx)
+                          for img_id in imds_ememplars]
+        assert(len(Tnovel) == nTestNovel)
+        assert(len(Exemplars) == len(Knovel) * nExemplars)
+        # random.shuffle(Exemplars)
+
+        return Tnovel, Exemplars
+
+    def sample_episode(self):
+        """Samples a training episode."""
+        nKnovel = self.nKnovel
+        nKbase = self.nKbase
+        nTestNovel = self.nTestNovel
+        nTestBase = self.nTestBase
+        nExemplars = self.nExemplars
+
+        Kbase, Knovel = self.sample_base_and_novel_categories(nKbase, nKnovel)
+        Tbase = self.sample_test_examples_for_base_categories(Kbase, nTestBase)
+        Tnovel, Exemplars = self.sample_train_and_test_examples_for_novel_categories(
+            Knovel, nTestNovel, nExemplars, nKbase)
+
+        # concatenate the base and novel category examples.
+        Test = Tbase + Tnovel
+        # random.shuffle(Test)
+        Kall = Kbase + Knovel
+
+        return Exemplars, Test, Kall, nKbase
+
+    def createExamplesTensorData(self, examples):
+        """
+        Creates the examples image and label tensor data.
+
+        Args:
+            examples: a list of 2-element tuples, each representing a
+                train or test example. The 1st element of each tuple
+                is the image id of the example and 2nd element is the
+                category label of the example, which is in the range
+                [0, nK - 1], where nK is the total number of categories
+                (both novel and base).
+
+        Returns:
+            images: a tensor of shape [nExamples, Height, Width, 3] with the
+                example images, where nExamples is the number of examples
+                (i.e., nExamples = len(examples)).
+            labels: a tensor of shape [nExamples] with the category label
+                of each example.
+        """
+        images = torch.stack(
+            [self.dataset[img_idx][0] for img_idx, _ in examples], dim=0)
+        labels = torch.LongTensor([label for _, label in examples])
+        return images, labels
+
+    def get_iterator(self, epoch=0):
+        rand_seed = epoch
+        random.seed(rand_seed)
+        np.random.seed(rand_seed)
+
+        def load_function(iter_idx):
+            Exemplars, Test, Kall, nKbase = self.sample_episode()
+            Xt, Yt = self.createExamplesTensorData(Test)
+            Kall = torch.LongTensor(Kall)
+            if len(Exemplars) > 0:
+                Xe, Ye = self.createExamplesTensorData(Exemplars)
+                return Xe, Ye, Xt, Yt, Kall, nKbase
+            else:
+                return Xt, Yt, Kall, nKbase
+
+        tnt_dataset = tnt.dataset.ListDataset(
+            elem_list=range(self.epoch_size), load=load_function)
+        data_loader = tnt_dataset.parallel(
+            batch_size=self.batch_size,
+            num_workers=(0 if self.is_eval_mode else self.num_workers),
+            shuffle=(False if self.is_eval_mode else True),)
+
+        return data_loader
+
+    def __call__(self, epoch=0):
+        return self.get_iterator(epoch)
+
+    def __len__(self):
+        return int(self.epoch_size / self.batch_size)

dataloader/chest1.py

@@ -0,0 +1,517 @@
+# Dataloader of Gidaris & Komodakis, CVPR 2018
+# Adapted from:
+# https://github.com/gidariss/FewShotWithoutForgetting/blob/master/dataloader.py
+from __future__ import print_function
+
+import os
+import os.path
+import numpy as npw
+import random
+import pickle
+import json
+import math
+
+import torch
+import torch.utils.data as data
+import torchvision
+import torchvision.datasets as datasets
+import torchvision.transforms as transforms
+import torchnet as tnt
+import numpy as np
+import pandas as pd
+
+
+import h5py
+
+import cv2
+from PIL import Image
+from PIL import ImageEnhance
+import matplotlib.pyplot as plt
+
+
+from torchvision.transforms.transforms import ToPILImage
+
+
+# Set the appropriate paths of the datasets here.
+# _CIFAR_FS_DATASET_DIR = './cifar/CIFAR-FS/'
+_CHEST_DATASET_DIR = './NIH'
+image_path = './NIH/images'
+
+
+label_dict = {'Cardiomegaly': 0, 'Edema': 1, 'Effusion': 2, 'Emphysema': 3, 'Infiltration': 4, 'Mass': 5, 'Atelectasis': 6, 'Consolidation': 7,
+              'Pleural_Thickening': 8, 'Fibrosis': 9, 'Hernia': 10, 'Pneumonia': 11, 'Nodule': 12, 'Pneumothorax': 13, 'No Finding': 14}
+
+
+def buildLabelIndex(labels):
+    label2inds = {}
+    for idx, label in enumerate(labels):
+        label = label_dict[label]
+        if label not in label2inds:
+            label2inds[label] = []
+        label2inds[label].append(idx)
+
+    return label2inds
+
+
+def load_data(file):
+    try:
+        with open(file, 'rb') as fo:
+            data = pickle.load(fo)
+        return data
+    except:
+        with open(file, 'rb') as f:
+            u = pickle._Unpickler(f)
+            u.encoding = 'latin1'
+            data = u.load()
+        return data
+
+
+class Chest(data.Dataset):
+    def __init__(self, phase='train', do_not_use_random_transf=False):
+
+        assert(phase == 'train' or phase == 'val' or phase ==
+               'test' or phase == 'trainval')
+        self.phase = phase
+        # self.name = phase + '.csv'
+
+        idx = 1  # represents group for experimentation
+
+        print('Loading Chest-XRay dataset - phase {0}'.format(phase))
+
+        train_path = os.path.join(_CHEST_DATASET_DIR, f'train{idx}.csv')
+        val_path = os.path.join(_CHEST_DATASET_DIR, f'val{idx}.csv')
+        test_path = os.path.join(_CHEST_DATASET_DIR, f'test{idx}.csv')
+
+        if self.phase == 'train':
+            # # During training phase we only load the training phase images
+            # # of the training categories (aka base categories).
+            # data_train = load_data(file_train_categories_train_phase)
+            # # self.data = data_train['data']
+            # self.labels = data_train['labels']
+
+            file = pd.read_csv(train_path)
+
+            self.data = file['image_id'].values
+
+            self.labels = file['class_name'].values
+
+            self.label2ind = buildLabelIndex(self.labels)
+
+            self.labelIds = sorted(self.label2ind.keys())
+            self.num_cats = len(self.labelIds)
+            self.labelIds_base = self.labelIds
+            self.num_cats_base = len(self.labelIds_base)
+
+        # elif self.phase == 'trainval':
+        #     # During training phase we only load the training phase images
+        #     # of the training categories (aka base categories).
+        #     data_train = load_data(file_train_categories_train_phase)
+        #     self.data = data_train['data']
+        #     self.labels = data_train['labels']
+        #     data_base = load_data(file_train_categories_val_phase)
+        #     data_novel = load_data(file_val_categories_val_phase)
+        #     self.data = np.concatenate(
+        #         [self.data, data_novel['data']], axis=0)
+        #     self.data = np.concatenate(
+        #         [self.data, data_base['data']], axis=0)
+
+        #     self.labels = np.concatenate(
+        #         [self.labels, data_novel['labels']], axis=0)
+        #     self.labels = np.concatenate(
+        #         [self.labels, data_base['labels']], axis=0)
+
+        #     self.label2ind = buildLabelIndex(self.labels)
+        #     self.labelIds = sorted(self.label2ind.keys())
+        #     self.num_cats = len(self.labelIds)
+        #     self.labelIds_base = self.labelIds
+        #     self.num_cats_base = len(self.labelIds_base)
+
+        elif self.phase == 'val' or self.phase == 'test':
+            if self.phase == 'test':
+                # # load data that will be used for evaluating the recognition
+                # # accuracy of the base categories.
+                # data_base = load_data(file_train_categories_test_phase)
+                # # load data that will be use for evaluating the few-shot recogniton
+                # # accuracy on the novel categories.
+                # data_novel = load_data(file_test_categories_test_phase)
+
+                train_file = pd.read_csv(train_path)
+                file = pd.read_csv(test_path)
+            else:  # phase=='val'
+                # # load data that will be used for evaluating the recognition
+                # # accuracy of the base categories.
+                # data_base = load_data(file_train_categories_val_phase)
+                # # load data that will be use for evaluating the few-shot recogniton
+                # # accuracy on the novel categories.
+                # data_novel = load_data(file_val_categories_val_phase)
+
+                train_file = pd.read_csv(train_path)
+                file = pd.read_csv(val_path)
+
+            # self.data = np.concatenate(
+            #     [data_base['data'], data_novel['data']], axis=0)
+            # self.labels = data_base['labels'] + data_novel['labels']
+
+            train_labels = train_file['class_name'].values
+            novel_labels = file['class_name'].values
+
+            self.data = np.concatenate(
+                [train_file['image_id'].values, file['image_id'].values], axis=0)
+            self.labels = np.concatenate(
+                [train_file['class_name'].values, file['class_name'].values], axis=0)
+
+
+            self.label2ind = buildLabelIndex(self.labels)
+            self.labelIds = sorted(self.label2ind.keys())
+            self.num_cats = len(self.labelIds)
+
+            # self.labelIds_base = buildLabelIndex(data_base['labels']).keys()
+            # self.labelIds_novel = buildLabelIndex(data_novel['labels']).keys()
+
+            self.labelIds_base = buildLabelIndex(train_labels).keys()
+            self.labelIds_novel = buildLabelIndex(novel_labels).keys()
+            print('='*60)
+            print(self.labelIds_novel)
+            print('='*60)
+
+            self.num_cats_base = len(self.labelIds_base)
+            self.num_cats_novel = len(self.labelIds_novel)
+            # print(self.labelIds_novel)
+            # print(self.num_cats_novel)
+            intersection = set(self.labelIds_base) & set(self.labelIds_novel)
+            assert(len(intersection) == 0)
+        else:
+            raise ValueError('Not valid phase {0}'.format(self.phase))
+
+        # mean_pix = [x/255.0 for x in [129.37731888,
+        #                               124.10583864, 112.47758569]]
+
+        # std_pix = [x/255.0 for x in [68.20947949, 65.43124043, 70.45866994]]
+
+        mean_pix =  [0.52024849, 0.52024849, 0.52024849]
+        std_pix =  [0.22699496, 0.22699496, 0.22699496]
+
+
+        normalize = transforms.Normalize(mean=mean_pix, std=std_pix)
+
+        if (self.phase == 'test' or self.phase == 'val') or (do_not_use_random_transf == True):
+
+            self.transform = transforms.Compose([
+                transforms.ToPILImage(),
+                # lambda x: np.asarray(x),
+                transforms.ToTensor(),
+                # lambda x: x/255.0,
+                normalize
+            ])
+        else:
+            self.transform = transforms.Compose([
+                transforms.ToPILImage(),
+                # transforms.RandomCrop(32, padding=4),
+                # transforms.ColorJitter(
+                #     brightness=0.4, contrast=0.4, saturation=0.4),
+                transforms.RandomHorizontalFlip(),
+                transforms.ToTensor(),
+                # lambda x: np.asarray(x),
+                # lambda x: x/255.0,
+                normalize
+            ])
+
+    def __getitem__(self, index):
+        img, label = cv2.imread(os.path.join(
+            image_path, self.data[index]))[:,:,::-1], self.labels[index]
+        img = cv2.resize(img,(128,128)) # resize by Garvit
+        # img = cv2.resize(img,(84, 84)) # resize by kshitiz
+
+        # img = Image.fromarray(img)
+        if self.transform is not None:
+            img = self.transform(img)
+        return img, label, self.data[index]
+        # return img, label
+
+    def __len__(self):
+        return len(self.data)
+
+
+class FewShotDataloader():
+    def __init__(self,
+                 dataset,
+                 nKnovel=5,  # number of novel categories.
+                 nKbase=-1,  # number of base categories.
+                 # number of training examples per novel category.
+                 nExemplars=1,
+                 # number of test examples for all the novel categories.
+                 nTestNovel=15*5,
+                 # number of test examples for all the base categories.
+                 nTestBase=15*5,
+                 batch_size=1,  # number of training episodes per batch.
+                 num_workers=4,
+                 epoch_size=2000,  # number of batches per epoch.
+                 ):
+
+        self.dataset = dataset
+        self.phase = self.dataset.phase
+        max_possible_nKnovel = (self.dataset.num_cats_base if self.phase == 'train' or self.phase == 'trainval'
+                                else self.dataset.num_cats_novel)
+
+        assert(nKnovel >= 0 and nKnovel <= max_possible_nKnovel)
+        self.nKnovel = nKnovel
+
+        max_possible_nKbase = self.dataset.num_cats_base
+        nKbase = nKbase if nKbase >= 0 else max_possible_nKbase
+        if (self.phase == 'train' or self.phase == 'trainval') and nKbase > 0:
+            nKbase -= self.nKnovel
+            max_possible_nKbase -= self.nKnovel
+
+        assert(nKbase >= 0 and nKbase <= max_possible_nKbase)
+        self.nKbase = nKbase
+
+        self.nExemplars = nExemplars
+        self.nTestNovel = nTestNovel
+        self.nTestBase = nTestBase
+        self.batch_size = batch_size
+        self.epoch_size = epoch_size
+        self.num_workers = num_workers
+        self.is_eval_mode = (self.phase == 'test') or (self.phase == 'val')
+
+    def sampleImageIdsFrom(self, cat_id, sample_size=1):
+        """
+        Samples `sample_size` number of unique image ids picked from the
+        category `cat_id` (i.e., self.dataset.label2ind[cat_id]).
+
+        Args:
+            cat_id: a scalar with the id of the category from which images will
+                be sampled.
+            sample_size: number of images that will be sampled.
+
+        Returns:
+            image_ids: a list of length `sample_size` with unique image ids.
+        """
+        assert(cat_id in self.dataset.label2ind)
+        assert(len(self.dataset.label2ind[cat_id]) >= sample_size)
+        # Note: random.sample samples elements without replacement.
+        # seed = random.randint(1,10000000)
+        # random.seed(seed)
+        return random.sample(self.dataset.label2ind[cat_id], sample_size)
+
+    def sampleCategories(self, cat_set, sample_size=1):
+        """
+        Samples `sample_size` number of unique categories picked from the
+        `cat_set` set of categories. `cat_set` can be either 'base' or 'novel'.
+
+        Args:
+            cat_set: string that specifies the set of categories from which
+                categories will be sampled.
+            sample_size: number of categories that will be sampled.
+
+        Returns:
+            cat_ids: a list of length `sample_size` with unique category ids.
+        """
+        if cat_set == 'base':
+            labelIds = self.dataset.labelIds_base
+        elif cat_set == 'novel':
+            labelIds = self.dataset.labelIds_novel
+        else:
+            raise ValueError('Not recognized category set {}'.format(cat_set))
+
+        assert(len(labelIds) >= sample_size)
+        # return sample_size unique categories chosen from labelIds set of
+        # categories (that can be either self.labelIds_base or self.labelIds_novel)
+        # Note: random.sample samples elements without replacement.
+        return random.sample(labelIds, sample_size)
+
+    def sample_base_and_novel_categories(self, nKbase, nKnovel):
+        """
+        Samples `nKbase` number of base categories and `nKnovel` number of novel
+        categories.
+
+        Args:
+            nKbase: number of base categories
+            nKnovel: number of novel categories
+
+        Returns:
+            Kbase: a list of length 'nKbase' with the ids of the sampled base
+                categories.
+            Knovel: a list of lenght 'nKnovel' with the ids of the sampled novel
+                categories.
+        """
+        if self.is_eval_mode:
+            assert(nKnovel <= self.dataset.num_cats_novel)
+            # sample from the set of base categories 'nKbase' number of base
+            # categories.
+            Kbase = sorted(self.sampleCategories('base', nKbase))
+            # sample from the set of novel categories 'nKnovel' number of novel
+            # categories.
+            Knovel = sorted(self.sampleCategories('novel', nKnovel))
+        else:
+            # sample from the set of base categories 'nKnovel' + 'nKbase' number
+            # of categories.
+            cats_ids = self.sampleCategories('base', nKnovel+nKbase)
+            assert(len(cats_ids) == (nKnovel+nKbase))
+            # Randomly pick 'nKnovel' number of fake novel categories and keep
+            # the rest as base categories.
+            random.shuffle(cats_ids)
+            Knovel = sorted(cats_ids[:nKnovel])
+            Kbase = sorted(cats_ids[nKnovel:])
+
+
+        return Kbase, Knovel
+
+    def sample_test_examples_for_base_categories(self, Kbase, nTestBase):
+        """
+        Sample `nTestBase` number of images from the `Kbase` categories.
+
+        Args:
+            Kbase: a list of length `nKbase` with the ids of the categories from
+                where the images will be sampled.
+            nTestBase: the total number of images that will be sampled.
+
+        Returns:
+            Tbase: a list of length `nTestBase` with 2-element tuples. The 1st
+                element of each tuple is the image id that was sampled and the
+                2nd elemend is its category label (which is in the range
+                [0, len(Kbase)-1]).
+        """
+        Tbase = []
+        if len(Kbase) > 0:
+            # Sample for each base category a number images such that the total
+            # number sampled images of all categories to be equal to `nTestBase`.
+            KbaseIndices = np.random.choice(
+                np.arange(len(Kbase)), size=nTestBase, replace=True)
+            KbaseIndices, NumImagesPerCategory = np.unique(
+                KbaseIndices, return_counts=True)
+
+            for Kbase_idx, NumImages in zip(KbaseIndices, NumImagesPerCategory):
+                imd_ids = self.sampleImageIdsFrom(
+                    Kbase[Kbase_idx], sample_size=NumImages)
+                Tbase += [(img_id, Kbase_idx) for img_id in imd_ids]
+
+        assert(len(Tbase) == nTestBase)
+
+        return Tbase
+
+    def sample_train_and_test_examples_for_novel_categories(
+            self, Knovel, nTestNovel, nExemplars, nKbase):
+        """Samples train and test examples of the novel categories.
+
+        Args:
+            Knovel: a list with the ids of the novel categories.
+            nTestNovel: the total number of test images that will be sampled
+                from all the novel categories.
+            nExemplars: the number of training examples per novel category that
+                will be sampled.
+            nKbase: the number of base categories. It is used as offset of the
+                category index of each sampled image.
+
+        Returns:
+            Tnovel: a list of length `nTestNovel` with 2-element tuples. The
+                1st element of each tuple is the image id that was sampled and
+                the 2nd element is its category label (which is in the range
+                [nKbase, nKbase + len(Knovel) - 1]).
+            Exemplars: a list of length len(Knovel) * nExemplars of 2-element
+                tuples. The 1st element of each tuple is the image id that was
+                sampled and the 2nd element is its category label (which is in
+                the ragne [nKbase, nKbase + len(Knovel) - 1]).
+        """
+
+        if len(Knovel) == 0:
+            return [], []
+
+        nKnovel = len(Knovel)
+        Tnovel = []
+        Exemplars = []
+        assert((nTestNovel % nKnovel) == 0)
+        nEvalExamplesPerClass = int(nTestNovel / nKnovel)
+
+        for Knovel_idx in range(len(Knovel)):
+            imd_ids = self.sampleImageIdsFrom(
+                Knovel[Knovel_idx],
+                sample_size=(nEvalExamplesPerClass + nExemplars))
+
+            imds_tnovel = imd_ids[:nEvalExamplesPerClass]
+            imds_ememplars = imd_ids[nEvalExamplesPerClass:]
+
+            Tnovel += [(img_id, nKbase+Knovel_idx) for img_id in imds_tnovel]
+            Exemplars += [(img_id, nKbase+Knovel_idx)
+                          for img_id in imds_ememplars]
+        assert(len(Tnovel) == nTestNovel)
+        assert(len(Exemplars) == len(Knovel) * nExemplars)
+        # random.shuffle(Exemplars)
+
+        return Tnovel, Exemplars
+
+    def sample_episode(self):
+        """Samples a training episode."""
+        nKnovel = self.nKnovel
+        nKbase = self.nKbase
+        nTestNovel = self.nTestNovel
+        nTestBase = self.nTestBase
+        nExemplars = self.nExemplars
+
+        Kbase, Knovel = self.sample_base_and_novel_categories(nKbase, nKnovel)
+        Tbase = self.sample_test_examples_for_base_categories(Kbase, nTestBase)
+        Tnovel, Exemplars = self.sample_train_and_test_examples_for_novel_categories(
+            Knovel, nTestNovel, nExemplars, nKbase)
+
+        # concatenate the base and novel category examples.
+        Test = Tbase + Tnovel
+        # random.shuffle(Test)
+        Kall = Kbase + Knovel
+
+        return Exemplars, Test, Kall, nKbase
+
+    def createExamplesTensorData(self, examples):
+        """
+        Creates the examples image and label tensor data.
+
+        Args:
+            examples: a list of 2-element tuples, each representing a
+                train or test example. The 1st element of each tuple
+                is the image id of the example and 2nd element is the
+                category label of the example, which is in the range
+                [0, nK - 1], where nK is the total number of categories
+                (both novel and base).
+
+        Returns:
+            images: a tensor of shape [nExamples, Height, Width, 3] with the
+                example images, where nExamples is the number of examples
+                (i.e., nExamples = len(examples)).
+            labels: a tensor of shape [nExamples] with the category label
+                of each example.
+        """
+        images = torch.stack(
+            [self.dataset[img_idx][0] for img_idx, _ in examples], dim=0)
+        names = np.stack(
+            [self.dataset[img_idx][-1] for img_idx, _ in examples], axis=0)
+        print(names)
+        labels = torch.LongTensor([label for _, label in examples])
+        return images, labels
+
+    def get_iterator(self, epoch=0):
+        rand_seed = epoch
+        random.seed(rand_seed)
+        np.random.seed(rand_seed)
+
+        def load_function(iter_idx):
+            Exemplars, Test, Kall, nKbase = self.sample_episode()
+            Xt, Yt = self.createExamplesTensorData(Test)
+            Kall = torch.LongTensor(Kall)
+            if len(Exemplars) > 0:
+                Xe, Ye = self.createExamplesTensorData(Exemplars)
+                return Xe, Ye, Xt, Yt, Kall, nKbase
+            else:
+                return Xt, Yt, Kall, nKbase
+
+        tnt_dataset = tnt.dataset.ListDataset(
+            elem_list=range(self.epoch_size), load=load_function)
+        data_loader = tnt_dataset.parallel(
+            batch_size=self.batch_size,
+            num_workers=(0 if self.is_eval_mode else self.num_workers),
+            shuffle=(False if self.is_eval_mode else True),)
+
+        return data_loader
+
+    def __call__(self, epoch=0):
+        return self.get_iterator(epoch)
+
+    def __len__(self):
+        return int(self.epoch_size / self.batch_size)

dataloader/mini_imagenet.py

@@ -0,0 +1,454 @@
+# Dataloader of Gidaris & Komodakis, CVPR 2018
+# Adapted from:
+# https://github.com/gidariss/FewShotWithoutForgetting/blob/master/dataloader.py
+from __future__ import print_function
+
+import os
+import os.path
+import numpy as np
+import random
+import pickle
+import json
+import math
+
+import torch
+import torch.utils.data as data
+import torchvision
+import torchvision.datasets as datasets
+import torchvision.transforms as transforms
+import torchnet as tnt
+
+import h5py
+
+from PIL import Image
+from PIL import ImageEnhance
+
+from pdb import set_trace as breakpoint
+
+from torchvision.transforms.transforms import ToPILImage
+
+# Set the appropriate paths of the datasets here.
+_MINI_IMAGENET_DATASET_DIR = './miniimagenet/' ## your miniimagenet folder
+
+
+def buildLabelIndex(labels):
+    label2inds = {}
+    for idx, label in enumerate(labels):
+        if label not in label2inds:
+            label2inds[label] = []
+        label2inds[label].append(idx)
+
+    return label2inds
+
+
+def load_data(file):
+    try:
+        with open(file, 'rb') as fo:
+            data = pickle.load(fo)
+        return data
+    except:
+        with open(file, 'rb') as f:
+            u = pickle._Unpickler(f)
+            u.encoding = 'latin1'
+            data = u.load()
+        return data
+
+class MiniImageNet(data.Dataset):
+    def __init__(self, phase='train', do_not_use_random_transf=False):
+
+        self.base_folder = 'miniImagenet'
+        #assert(phase=='train' or phase=='val' or phase=='test' or ph)
+        self.phase = phase
+        self.name = 'MiniImageNet_' + phase
+
+        print('Loading mini ImageNet dataset - phase {0}'.format(phase))
+        file_train_categories_train_phase = os.path.join(
+            _MINI_IMAGENET_DATASET_DIR,
+            'miniImageNet_category_split_train_phase_train.pickle')
+        file_train_categories_val_phase = os.path.join(
+            _MINI_IMAGENET_DATASET_DIR,
+            'miniImageNet_category_split_train_phase_val.pickle')
+        file_train_categories_test_phase = os.path.join(
+            _MINI_IMAGENET_DATASET_DIR,
+            'miniImageNet_category_split_train_phase_test.pickle')
+        file_val_categories_val_phase = os.path.join(
+            _MINI_IMAGENET_DATASET_DIR,
+            'miniImageNet_category_split_val.pickle')
+        file_test_categories_test_phase = os.path.join(
+            _MINI_IMAGENET_DATASET_DIR,
+            'miniImageNet_category_split_test.pickle')
+
+        if self.phase=='train':
+            # During training phase we only load the training phase images
+            # of the training categories (aka base categories).
+            data_train = load_data(file_train_categories_train_phase)
+            self.data = data_train['data']
+            self.labels = data_train['labels']
+
+            self.label2ind = buildLabelIndex(self.labels)
+            self.labelIds = sorted(self.label2ind.keys())
+            self.num_cats = len(self.labelIds)
+            self.labelIds_base = self.labelIds
+            self.num_cats_base = len(self.labelIds_base)
+        elif self.phase == 'trainval':
+            # During training phase we only load the training phase images
+            # of the training categories (aka base categories).
+            data_train = load_data(file_train_categories_train_phase)
+            self.data = data_train['data']
+            self.labels = data_train['labels']
+            data_base = load_data(file_train_categories_val_phase)
+            data_novel = load_data(file_val_categories_val_phase)
+            self.data = np.concatenate(
+                [self.data, data_novel['data']], axis=0)
+            self.data = np.concatenate(
+                [self.data, data_base['data']], axis=0)
+            self.labels = np.concatenate(
+                [self.labels, data_novel['labels']], axis=0)
+            self.labels = np.concatenate(
+                [self.labels, data_base['labels']], axis=0)
+
+            self.label2ind = buildLabelIndex(self.labels)
+            self.labelIds = sorted(self.label2ind.keys())
+            self.num_cats = len(self.labelIds)
+            self.labelIds_base = self.labelIds
+            self.num_cats_base = len(self.labelIds_base)
+        elif self.phase=='val' or self.phase=='test':
+            if self.phase=='test':
+                # load data that will be used for evaluating the recognition
+                # accuracy of the base categories.
+                data_base = load_data(file_train_categories_test_phase)
+                # load data that will be use for evaluating the few-shot recogniton
+                # accuracy on the novel categories.
+                data_novel = load_data(file_test_categories_test_phase)
+            else: # phase=='val'
+                # load data that will be used for evaluating the recognition
+                # accuracy of the base categories.
+                data_base = load_data(file_train_categories_val_phase)
+                # load data that will be use for evaluating the few-shot recogniton
+                # accuracy on the novel categories.
+                data_novel = load_data(file_val_categories_val_phase)
+
+            self.data = np.concatenate(
+                [data_base['data'], data_novel['data']], axis=0)
+            self.labels = data_base['labels'] + data_novel['labels']
+
+            self.label2ind = buildLabelIndex(self.labels)
+            self.labelIds = sorted(self.label2ind.keys())
+            self.num_cats = len(self.labelIds)
+
+            self.labelIds_base = buildLabelIndex(data_base['labels']).keys()
+            self.labelIds_novel = buildLabelIndex(data_novel['labels']).keys()
+            self.num_cats_base = len(self.labelIds_base)
+            self.num_cats_novel = len(self.labelIds_novel)
+            intersection = set(self.labelIds_base) & set(self.labelIds_novel)
+            assert(len(intersection) == 0)
+        else:
+            raise ValueError('Not valid phase {0}'.format(self.phase))
+
+        mean_pix = [x/255.0 for x in [120.39586422,  115.59361427, 104.54012653]]
+        std_pix = [x/255.0 for x in [70.68188272,  68.27635443,  72.54505529]]
+        normalize = transforms.Normalize(mean=mean_pix, std=std_pix)
+
+        if (self.phase=='test' or self.phase=='val') or (do_not_use_random_transf==True):
+            self.transform = transforms.Compose([
+                # transforms.ToPILImage(),
+                # lambda x: np.asarray(x),
+                transforms.ToTensor(),
+                normalize
+            ])
+        else:
+            self.transform = transforms.Compose([
+                # transforms.ToPILImage(),
+                transforms.RandomCrop(84, padding=8),
+                transforms.ColorJitter(brightness=0.4, contrast=0.4, saturation=0.4),
+                transforms.RandomHorizontalFlip(),
+                # lambda x: np.asarray(x),
+                transforms.ToTensor(),
+                normalize
+            ])
+            
+    def __getitem__(self, index):
+        img, label = self.data[index], self.labels[index]
+        # doing this so that it is consistent with all other datasets
+        # to return a PIL Image
+        img = Image.fromarray(img)
+        if self.transform is not None:
+            img = self.transform(img)
+        return img, label
+
+    def __len__(self):
+        return len(self.data)
+
+
+class FewShotDataloader():
+    def __init__(self,
+                 dataset,
+                 nKnovel=5, # number of novel categories.
+                 nKbase=-1, # number of base categories.
+                 nExemplars=1, # number of training examples per novel category.
+                 nTestNovel=15*5, # number of test examples for all the novel categories.
+                 nTestBase=15*5, # number of test examples for all the base categories.
+                 batch_size=1, # number of training episodes per batch.
+                 num_workers=0,
+                 epoch_size=2000, # number of batches per epoch.
+                 ):
+
+        self.dataset = dataset
+        self.phase = self.dataset.phase
+        max_possible_nKnovel = (self.dataset.num_cats_base if self.phase=='train' or self.phase=='trainval'
+                                else self.dataset.num_cats_novel)
+        assert(nKnovel >= 0 and nKnovel < max_possible_nKnovel)
+        self.nKnovel = nKnovel
+
+        max_possible_nKbase = self.dataset.num_cats_base
+        nKbase = nKbase if nKbase >= 0 else max_possible_nKbase
+        if (self.phase=='train'or self.phase=='trainval') and nKbase > 0:
+            nKbase -= self.nKnovel
+            max_possible_nKbase -= self.nKnovel
+
+        assert(nKbase >= 0 and nKbase <= max_possible_nKbase)
+        self.nKbase = nKbase
+
+        self.nExemplars = nExemplars
+        self.nTestNovel = nTestNovel
+        self.nTestBase = nTestBase
+        self.batch_size = batch_size
+        self.epoch_size = epoch_size
+        self.num_workers = num_workers
+        self.is_eval_mode = (self.phase=='test') or (self.phase=='val')
+
+    def sampleImageIdsFrom(self, cat_id, sample_size=1):
+        """
+        Samples `sample_size` number of unique image ids picked from the
+        category `cat_id` (i.e., self.dataset.label2ind[cat_id]).
+
+        Args:
+            cat_id: a scalar with the id of the category from which images will
+                be sampled.
+            sample_size: number of images that will be sampled.
+
+        Returns:
+            image_ids: a list of length `sample_size` with unique image ids.
+        """
+        assert(cat_id in self.dataset.label2ind)
+        assert(len(self.dataset.label2ind[cat_id]) >= sample_size)
+        # Note: random.sample samples elements without replacement.
+        return random.sample(self.dataset.label2ind[cat_id], sample_size)
+
+    def sampleCategories(self, cat_set, sample_size=1):
+        """
+        Samples `sample_size` number of unique categories picked from the
+        `cat_set` set of categories. `cat_set` can be either 'base' or 'novel'.
+
+        Args:
+            cat_set: string that specifies the set of categories from which
+                categories will be sampled.
+            sample_size: number of categories that will be sampled.
+
+        Returns:
+            cat_ids: a list of length `sample_size` with unique category ids.
+        """
+        if cat_set=='base':
+            labelIds = self.dataset.labelIds_base
+        elif cat_set=='novel':
+            labelIds = self.dataset.labelIds_novel
+        else:
+            raise ValueError('Not recognized category set {}'.format(cat_set))
+
+        assert(len(labelIds) >= sample_size)
+        # return sample_size unique categories chosen from labelIds set of
+        # categories (that can be either self.labelIds_base or self.labelIds_novel)
+        # Note: random.sample samples elements without replacement.
+        return random.sample(labelIds, sample_size)
+
+    def sample_base_and_novel_categories(self, nKbase, nKnovel):
+        """
+        Samples `nKbase` number of base categories and `nKnovel` number of novel
+        categories.
+
+        Args:
+            nKbase: number of base categories
+            nKnovel: number of novel categories
+
+        Returns:
+            Kbase: a list of length 'nKbase' with the ids of the sampled base
+                categories.
+            Knovel: a list of lenght 'nKnovel' with the ids of the sampled novel
+                categories.
+        """
+        if self.is_eval_mode:
+            assert(nKnovel <= self.dataset.num_cats_novel)
+            # sample from the set of base categories 'nKbase' number of base
+            # categories.
+            Kbase = sorted(self.sampleCategories('base', nKbase))
+            # sample from the set of novel categories 'nKnovel' number of novel
+            # categories.
+            Knovel = sorted(self.sampleCategories('novel', nKnovel))
+        else:
+            # sample from the set of base categories 'nKnovel' + 'nKbase' number
+            # of categories.
+            cats_ids = self.sampleCategories('base', nKnovel+nKbase)
+            assert(len(cats_ids) == (nKnovel+nKbase))
+            # Randomly pick 'nKnovel' number of fake novel categories and keep
+            # the rest as base categories.
+            random.shuffle(cats_ids)
+            Knovel = sorted(cats_ids[:nKnovel])
+            Kbase = sorted(cats_ids[nKnovel:])
+
+        return Kbase, Knovel
+
+    def sample_test_examples_for_base_categories(self, Kbase, nTestBase):
+        """
+        Sample `nTestBase` number of images from the `Kbase` categories.
+
+        Args:
+            Kbase: a list of length `nKbase` with the ids of the categories from
+                where the images will be sampled.
+            nTestBase: the total number of images that will be sampled.
+
+        Returns:
+            Tbase: a list of length `nTestBase` with 2-element tuples. The 1st
+                element of each tuple is the image id that was sampled and the
+                2nd elemend is its category label (which is in the range
+                [0, len(Kbase)-1]).
+        """
+        Tbase = []
+        if len(Kbase) > 0:
+            # Sample for each base category a number images such that the total
+            # number sampled images of all categories to be equal to `nTestBase`.
+            KbaseIndices = np.random.choice(
+                np.arange(len(Kbase)), size=nTestBase, replace=True)
+            KbaseIndices, NumImagesPerCategory = np.unique(
+                KbaseIndices, return_counts=True)
+
+            for Kbase_idx, NumImages in zip(KbaseIndices, NumImagesPerCategory):
+                imd_ids = self.sampleImageIdsFrom(
+                    Kbase[Kbase_idx], sample_size=NumImages)
+                Tbase += [(img_id, Kbase_idx) for img_id in imd_ids]
+
+        assert(len(Tbase) == nTestBase)
+
+        return Tbase
+
+    def sample_train_and_test_examples_for_novel_categories(
+            self, Knovel, nTestNovel, nExemplars, nKbase):
+        """Samples train and test examples of the novel categories.
+
+        Args:
+    	    Knovel: a list with the ids of the novel categories.
+            nTestNovel: the total number of test images that will be sampled
+                from all the novel categories.
+            nExemplars: the number of training examples per novel category that
+                will be sampled.
+            nKbase: the number of base categories. It is used as offset of the
+                category index of each sampled image.
+
+        Returns:
+            Tnovel: a list of length `nTestNovel` with 2-element tuples. The
+                1st element of each tuple is the image id that was sampled and
+                the 2nd element is its category label (which is in the range
+                [nKbase, nKbase + len(Knovel) - 1]).
+            Exemplars: a list of length len(Knovel) * nExemplars of 2-element
+                tuples. The 1st element of each tuple is the image id that was
+                sampled and the 2nd element is its category label (which is in
+                the ragne [nKbase, nKbase + len(Knovel) - 1]).
+        """
+
+        if len(Knovel) == 0:
+            return [], []
+
+        nKnovel = len(Knovel)
+        Tnovel = []
+        Exemplars = []
+        assert((nTestNovel % nKnovel) == 0)
+        nEvalExamplesPerClass = int(nTestNovel / nKnovel)
+
+        for Knovel_idx in range(nKnovel):
+            imd_ids = self.sampleImageIdsFrom(
+                Knovel[Knovel_idx],
+                sample_size=(nEvalExamplesPerClass + nExemplars))
+
+            imds_tnovel = imd_ids[:nEvalExamplesPerClass]
+            imds_ememplars = imd_ids[nEvalExamplesPerClass:]
+
+            Tnovel += [(img_id, nKbase+Knovel_idx) for img_id in imds_tnovel]
+            Exemplars += [(img_id, nKbase+Knovel_idx) for img_id in imds_ememplars]
+        assert(len(Tnovel) == nTestNovel)
+        assert(len(Exemplars) == len(Knovel) * nExemplars)
+
+        # random.shuffle(Exemplars)
+
+        return Tnovel, Exemplars
+
+    def sample_episode(self):
+        """Samples a training episode."""
+        nKnovel = self.nKnovel
+        nKbase = self.nKbase
+        nTestNovel = self.nTestNovel
+        nTestBase = self.nTestBase
+        nExemplars = self.nExemplars
+
+        Kbase, Knovel = self.sample_base_and_novel_categories(nKbase, nKnovel)
+        Tbase = self.sample_test_examples_for_base_categories(Kbase, nTestBase)
+        # print(Kbase,Knovel,Tbase)
+        Tnovel, Exemplars = self.sample_train_and_test_examples_for_novel_categories(
+            Knovel, nTestNovel, nExemplars, nKbase)
+        # concatenate the base and novel category examples.
+        Test = Tbase + Tnovel
+        # random.shuffle(Test)
+        Kall = Kbase + Knovel
+
+        return Exemplars, Test, Kall, nKbase
+
+    def createExamplesTensorData(self, examples):
+        """
+        Creates the examples image and label tensor data.
+
+        Args:
+            examples: a list of 2-element tuples, each representing a
+                train or test example. The 1st element of each tuple
+                is the image id of the example and 2nd element is the
+                category label of the example, which is in the range
+                [0, nK - 1], where nK is the total number of categories
+                (both novel and base).
+
+        Returns:
+            images: a tensor of shape [nExamples, Height, Width, 3] with the
+                example images, where nExamples is the number of examples
+                (i.e., nExamples = len(examples)).
+            labels: a tensor of shape [nExamples] with the category label
+                of each example.
+        """
+        images = torch.stack(
+            [self.dataset[img_idx][0] for img_idx, _ in examples], dim=0)
+        labels = torch.LongTensor([label for _, label in examples])
+        return images, labels
+
+    def get_iterator(self, epoch=0):
+        rand_seed = epoch
+        random.seed(rand_seed)
+        np.random.seed(rand_seed)
+        def load_function(iter_idx):
+            Exemplars, Test, Kall, nKbase = self.sample_episode()
+            Xt, Yt = self.createExamplesTensorData(Test)
+            Kall = torch.LongTensor(Kall)
+            if len(Exemplars) > 0:
+                Xe, Ye = self.createExamplesTensorData(Exemplars)
+                return Xe, Ye, Xt, Yt, Kall, nKbase
+            else:
+                return Xt, Yt, Kall, nKbase
+
+        tnt_dataset = tnt.dataset.ListDataset(
+            elem_list=range(self.epoch_size), load=load_function)
+        data_loader = tnt_dataset.parallel(
+            batch_size=self.batch_size,
+            num_workers=(0 if self.is_eval_mode else self.num_workers),
+            shuffle=(False if self.is_eval_mode else True))
+
+        return data_loader
+
+    def __call__(self, epoch=0):
+        return self.get_iterator(epoch)
+
+    def __len__(self):
+        return int(self.epoch_size / self.batch_size)

dataloader/simple_datamanager.py

@@ -0,0 +1,43 @@
+import torch
+from abc import abstractmethod
+import os
+from PIL import Image
+import json
+
+class DataManager:
+    @abstractmethod
+    def get_data_loader(self, data_file, aug):
+        pass
+
+
+class SimpleDataset:
+    def __init__(self, data_file, transform):
+        with open(data_file, 'r') as f:
+            self.meta = json.load(f)
+        self.transform = transform
+        #self.target_transform = target_transform
+
+
+    def __getitem__(self,i):
+        image_path = os.path.join(self.meta['image_names'][i])
+        img = Image.open(image_path).convert('RGB')
+        img = self.transform(img)
+        target = self.target_transform(self.meta['image_labels'][i])
+        return img, target
+
+    def __len__(self):
+        return len(self.meta['image_names'])
+
+
+class SimpleDataManager(DataManager):
+    def __init__(self, dataset, batch_size):
+        super(SimpleDataManager, self).__init__()
+        self.batch_size = batch_size
+        self.dataset = dataset
+
+    def get_data_loader(self): #parameters that would change on train/val set
+        dataset = self.dataset#SimpleDataset(data_file, transform)
+        data_loader_params = dict(batch_size = self.batch_size, shuffle = True, num_workers = 12, pin_memory = True)
+        data_loader = torch.utils.data.DataLoader(dataset, **data_loader_params)
+
+        return data_loader

dataloader/tiered_imagenet.py

@@ -0,0 +1,512 @@
+# Dataloader of Gidaris & Komodakis, CVPR 2018
+# Adapted from:
+# https://github.com/gidariss/FewShotWithoutForgetting/blob/master/dataloader.py
+from __future__ import print_function
+
+import os
+import os.path
+import numpy as np
+import random
+import pickle
+import json
+import math
+
+import torch
+import torch.utils.data as data
+import torchvision
+import torchvision.datasets as datasets
+import torchvision.transforms as transforms
+import torchnet as tnt
+
+import h5py
+
+from PIL import Image
+from PIL import ImageEnhance
+
+from pdb import set_trace as breakpoint
+
+from torchvision.transforms.transforms import ToPILImage
+
+# Set the appropriate paths of the datasets here.
+_TIERED_IMAGENET_DATASET_DIR = './tieredimagenet/' # your tiered imagenet folder
+
+def buildLabelIndex(labels):
+    label2inds = {}
+    for idx, label in enumerate(labels):
+        if label not in label2inds:
+            label2inds[label] = []
+        label2inds[label].append(idx)
+
+    return label2inds
+
+
+def load_data(file):
+    try:
+        with open(file, 'rb') as fo:
+            data = pickle.load(fo)
+        return data
+    except:
+        with open(file, 'rb') as f:
+            u = pickle._Unpickler(f)
+            u.encoding = 'latin1'
+            data = u.load()
+        return data
+
+class tieredImageNet(data.Dataset):
+    def __init__(self, phase='train', do_not_use_random_transf=False):
+
+        assert(phase=='train' or phase=='val' or phase=='test' or phase=='trainval')
+        self.phase = phase
+        self.name = 'tieredImageNet_' + phase
+
+        print('Loading tiered ImageNet dataset - phase {0}'.format(phase))
+        file_train_categories_train_phase = os.path.join(
+            _TIERED_IMAGENET_DATASET_DIR,
+            'train_images.npz')
+        label_train_categories_train_phase = os.path.join(
+            _TIERED_IMAGENET_DATASET_DIR,
+            'train_labels.pkl')
+        file_train_categories_val_phase = os.path.join(
+            _TIERED_IMAGENET_DATASET_DIR,
+            'train_images.npz')
+        label_train_categories_val_phase = os.path.join(
+            _TIERED_IMAGENET_DATASET_DIR,
+            'train_labels.pkl')
+        file_train_categories_test_phase = os.path.join(
+            _TIERED_IMAGENET_DATASET_DIR,
+            'train_images.npz')
+        label_train_categories_test_phase = os.path.join(
+            _TIERED_IMAGENET_DATASET_DIR,
+            'train_labels.pkl')
+
+        file_val_categories_val_phase = os.path.join(
+            _TIERED_IMAGENET_DATASET_DIR,
+            'val_images.npz')
+        label_val_categories_val_phase = os.path.join(
+            _TIERED_IMAGENET_DATASET_DIR,
+            'val_labels.pkl')
+        file_test_categories_test_phase = os.path.join(
+            _TIERED_IMAGENET_DATASET_DIR,
+            'test_images.npz')
+        label_test_categories_test_phase = os.path.join(
+            _TIERED_IMAGENET_DATASET_DIR,
+            'test_labels.pkl')
+
+        if self.phase == 'train':
+            # During training phase we only load the training phase images
+            # of the training categories (aka base categories).
+            data_train = load_data(label_train_categories_train_phase)
+            # self.data = data_train['data']
+            self.labels = data_train['labels']
+            self.data = np.load(file_train_categories_train_phase)[
+                'images']  # np.array(load_data(file_train_categories_train_phase))
+            # self.labels = load_data(file_train_categories_train_phase)#data_train['labels']
+
+            self.label2ind = buildLabelIndex(self.labels)
+            self.labelIds = sorted(self.label2ind.keys())
+            self.num_cats = len(self.labelIds)
+            self.labelIds_base = self.labelIds
+            self.num_cats_base = len(self.labelIds_base)
+        # if self.phase=='train':
+        #     # During training phase we only load the training phase images
+        #     # of the training categories (aka base categories).
+        #     data_train = load_data(label_train_categories_train_phase)
+        #     #self.data = data_train['data']
+        #     self.labels = data_train['labels']
+        #     self.data = np.load(file_train_categories_train_phase)['images']#np.array(load_data(file_train_categories_train_phase))
+        #     #self.labels = load_data(file_train_categories_train_phase)#data_train['labels']
+        #
+        #
+        #     data_base = load_data(label_train_categories_val_phase)['labels']
+        #     data_base_images = np.load(file_train_categories_val_phase)['images']
+        #     data_novel = load_data(label_val_categories_val_phase)['labels']
+        #     data_novel_images = np.load(file_val_categories_val_phase)['images']
+        #
+        #     self.data = np.concatenate(
+        #         [self.data, data_base_images], axis=0)
+        #     self.data = np.concatenate(
+        #         [self.data, data_novel_images], axis=0)
+        #     self.labels = np.concatenate(
+        #         [self.labels, data_base], axis=0)
+        #     self.labels = np.concatenate(
+        #         [self.labels, data_novel], axis=0)
+        #
+        #
+        #     self.label2ind = buildLabelIndex(self.labels)
+        #     self.labelIds = sorted(self.label2ind.keys())
+        #     self.num_cats = len(self.labelIds)
+        #     self.labelIds_base = self.labelIds
+        #     self.num_cats_base = len(self.labelIds_base)
+        elif self.phase == 'trainval':
+            # During training phase we only load the training phase images
+            # of the training categories (aka base categories).
+            data_train = load_data(file_train_categories_train_phase)
+            #self.data = data_train['data']
+            self.data = np.load(file_train_categories_train_phase)['images']
+            self.labels = data_train['labels']
+
+            data_base = load_data(label_train_categories_val_phase)['labels']
+            data_base_images = np.load(file_train_categories_val_phase)['images']
+            data_novel = load_data(label_val_categories_val_phase)['labels']
+            data_novel_images = np.load(file_val_categories_val_phase)['images']
+
+            self.data = np.concatenate(
+                [self.data, data_base_images], axis=0)
+            self.data = np.concatenate(
+                [self.data, data_novel_images], axis=0)
+            self.labels = np.concatenate(
+                [self.labels, data_base], axis=0)
+            self.labels = np.concatenate(
+                [self.labels, data_novel], axis=0)
+
+            self.label2ind = buildLabelIndex(self.labels)
+            self.labelIds = sorted(self.label2ind.keys())
+            self.num_cats = len(self.labelIds)
+            self.labelIds_base = self.labelIds
+            self.num_cats_base = len(self.labelIds_base)
+        elif self.phase=='val' or self.phase=='test':
+            if self.phase=='test':
+                # load data that will be used for evaluating the recognition
+                # accuracy of the base categories.
+                data_base = load_data(label_train_categories_test_phase)
+                data_base_images = np.load(file_train_categories_test_phase)['images']
+                
+                # load data that will be use for evaluating the few-shot recogniton
+                # accuracy on the novel categories.
+                data_novel = load_data(label_test_categories_test_phase)
+                data_novel_images = np.load(file_test_categories_test_phase)['images']
+            else: # phase=='val'
+                # load data that will be used for evaluating the recognition
+                # accuracy of the base categories.
+                data_base = load_data(label_train_categories_val_phase)
+                data_base_images = np.load(file_train_categories_val_phase)['images']
+                #print (data_base_images)
+                #print (data_base_images.shape)
+                # load data that will be use for evaluating the few-shot recogniton
+                # accuracy on the novel categories.
+                data_novel = load_data(label_val_categories_val_phase)
+                data_novel_images = np.load(file_val_categories_val_phase)['images']
+
+            self.data = np.concatenate(
+                [data_base_images, data_novel_images], axis=0)
+            self.labels = data_base['labels'] + data_novel['labels']
+
+            self.label2ind = buildLabelIndex(self.labels)
+            self.labelIds = sorted(self.label2ind.keys())
+            self.num_cats = len(self.labelIds)
+
+            self.labelIds_base = buildLabelIndex(data_base['labels']).keys()
+            self.labelIds_novel = buildLabelIndex(data_novel['labels']).keys()
+            self.num_cats_base = len(self.labelIds_base)
+            self.num_cats_novel = len(self.labelIds_novel)
+            intersection = set(self.labelIds_base) & set(self.labelIds_novel)
+            print (intersection)
+            assert(len(intersection) == 0)
+        else:
+            raise ValueError('Not valid phase {0}'.format(self.phase))
+
+        mean_pix = [x/255.0 for x in [120.39586422,  115.59361427, 104.54012653]]
+        std_pix = [x/255.0 for x in [70.68188272,  68.27635443,  72.54505529]]
+        normalize = transforms.Normalize(mean=mean_pix, std=std_pix)
+
+        if (self.phase=='test' or self.phase=='val') or (do_not_use_random_transf==True):
+            self.transform = transforms.Compose([
+                # lambda x: np.asarray(x),
+                transforms.ToTensor(),
+                normalize
+            ])
+        else:
+            self.transform = transforms.Compose([
+                transforms.RandomCrop(84, padding=8),
+                transforms.ColorJitter(brightness=0.4, contrast=0.4, saturation=0.4),
+                transforms.RandomHorizontalFlip(),
+                # lambda x: np.asarray(x),
+                transforms.ToTensor(),
+                normalize
+            ])
+
+    def __getitem__(self, index):
+        img, label = self.data[index], self.labels[index]
+        # doing this so that it is consistent with all other datasets
+        # to return a PIL Image
+        img = Image.fromarray(img)
+        if self.transform is not None:
+            img = self.transform(img)
+        return img, label
+
+    def __len__(self):
+        return len(self.data)
+
+
+class FewShotDataloader():
+    def __init__(self,
+                 dataset,
+                 nKnovel=5, # number of novel categories.
+                 nKbase=-1, # number of base categories.
+                 nExemplars=1, # number of training examples per novel category.
+                 nTestNovel=15*5, # number of test examples for all the novel categories.
+                 nTestBase=15*5, # number of test examples for all the base categories.
+                 batch_size=1, # number of training episodes per batch.
+                 num_workers=1,
+                 epoch_size=2000, # number of batches per epoch.
+                 ):
+
+        self.dataset = dataset
+        self.phase = self.dataset.phase
+        max_possible_nKnovel = (self.dataset.num_cats_base if self.phase=='train' or self.phase=='trainval'
+                                else self.dataset.num_cats_novel)
+        assert(nKnovel >= 0 and nKnovel < max_possible_nKnovel)
+        self.nKnovel = nKnovel
+
+        max_possible_nKbase = self.dataset.num_cats_base
+        nKbase = nKbase if nKbase >= 0 else max_possible_nKbase
+        if (self.phase=='train'or self.phase=='trainval') and nKbase > 0:
+            nKbase -= self.nKnovel
+            max_possible_nKbase -= self.nKnovel
+
+        assert(nKbase >= 0 and nKbase <= max_possible_nKbase)
+        self.nKbase = nKbase
+
+        self.nExemplars = nExemplars
+        self.nTestNovel = nTestNovel
+        self.nTestBase = nTestBase
+        self.batch_size = batch_size
+        self.epoch_size = epoch_size
+        self.num_workers = num_workers
+        self.is_eval_mode = (self.phase=='test') or (self.phase=='val')
+
+    def sampleImageIdsFrom(self, cat_id, sample_size=1):
+        """
+        Samples `sample_size` number of unique image ids picked from the
+        category `cat_id` (i.e., self.dataset.label2ind[cat_id]).
+
+        Args:
+            cat_id: a scalar with the id of the category from which images will
+                be sampled.
+            sample_size: number of images that will be sampled.
+
+        Returns:
+            image_ids: a list of length `sample_size` with unique image ids.
+        """
+        assert(cat_id in self.dataset.label2ind)
+        assert(len(self.dataset.label2ind[cat_id]) >= sample_size)
+        # Note: random.sample samples elements without replacement.
+        return random.sample(self.dataset.label2ind[cat_id], sample_size)
+
+    def sampleCategories(self, cat_set, sample_size=1):
+        """
+        Samples `sample_size` number of unique categories picked from the
+        `cat_set` set of categories. `cat_set` can be either 'base' or 'novel'.
+
+        Args:
+            cat_set: string that specifies the set of categories from which
+                categories will be sampled.
+            sample_size: number of categories that will be sampled.
+
+        Returns:
+            cat_ids: a list of length `sample_size` with unique category ids.
+        """
+        if cat_set=='base':
+            labelIds = self.dataset.labelIds_base
+        elif cat_set=='novel':
+            labelIds = self.dataset.labelIds_novel
+        else:
+            raise ValueError('Not recognized category set {}'.format(cat_set))
+
+        assert(len(labelIds) >= sample_size)
+        # return sample_size unique categories chosen from labelIds set of
+        # categories (that can be either self.labelIds_base or self.labelIds_novel)
+        # Note: random.sample samples elements without replacement.
+        return random.sample(labelIds, sample_size)
+
+    def sample_base_and_novel_categories(self, nKbase, nKnovel):
+        """
+        Samples `nKbase` number of base categories and `nKnovel` number of novel
+        categories.
+
+        Args:
+            nKbase: number of base categories
+            nKnovel: number of novel categories
+
+        Returns:
+            Kbase: a list of length 'nKbase' with the ids of the sampled base
+                categories.
+            Knovel: a list of lenght 'nKnovel' with the ids of the sampled novel
+                categories.
+        """
+        if self.is_eval_mode:
+            assert(nKnovel <= self.dataset.num_cats_novel)
+            # sample from the set of base categories 'nKbase' number of base
+            # categories.
+            Kbase = sorted(self.sampleCategories('base', nKbase))
+            # sample from the set of novel categories 'nKnovel' number of novel
+            # categories.
+            Knovel = sorted(self.sampleCategories('novel', nKnovel))
+        else:
+            # sample from the set of base categories 'nKnovel' + 'nKbase' number
+            # of categories.
+            cats_ids = self.sampleCategories('base', nKnovel+nKbase)
+            assert(len(cats_ids) == (nKnovel+nKbase))
+            # Randomly pick 'nKnovel' number of fake novel categories and keep
+            # the rest as base categories.
+            random.shuffle(cats_ids)
+            Knovel = sorted(cats_ids[:nKnovel])
+            Kbase = sorted(cats_ids[nKnovel:])
+
+        return Kbase, Knovel
+
+    def sample_test_examples_for_base_categories(self, Kbase, nTestBase):
+        """
+        Sample `nTestBase` number of images from the `Kbase` categories.
+
+        Args:
+            Kbase: a list of length `nKbase` with the ids of the categories from
+                where the images will be sampled.
+            nTestBase: the total number of images that will be sampled.
+
+        Returns:
+            Tbase: a list of length `nTestBase` with 2-element tuples. The 1st
+                element of each tuple is the image id that was sampled and the
+                2nd elemend is its category label (which is in the range
+                [0, len(Kbase)-1]).
+        """
+        Tbase = []
+        if len(Kbase) > 0:
+            # Sample for each base category a number images such that the total
+            # number sampled images of all categories to be equal to `nTestBase`.
+            KbaseIndices = np.random.choice(
+                np.arange(len(Kbase)), size=nTestBase, replace=True)
+            KbaseIndices, NumImagesPerCategory = np.unique(
+                KbaseIndices, return_counts=True)
+
+            for Kbase_idx, NumImages in zip(KbaseIndices, NumImagesPerCategory):
+                imd_ids = self.sampleImageIdsFrom(
+                    Kbase[Kbase_idx], sample_size=NumImages)
+                Tbase += [(img_id, Kbase_idx) for img_id in imd_ids]
+
+        assert(len(Tbase) == nTestBase)
+
+        return Tbase
+
+    def sample_train_and_test_examples_for_novel_categories(
+            self, Knovel, nTestNovel, nExemplars, nKbase):
+        """Samples train and test examples of the novel categories.
+
+        Args:
+    	    Knovel: a list with the ids of the novel categories.
+            nTestNovel: the total number of test images that will be sampled
+                from all the novel categories.
+            nExemplars: the number of training examples per novel category that
+                will be sampled.
+            nKbase: the number of base categories. It is used as offset of the
+                category index of each sampled image.
+
+        Returns:
+            Tnovel: a list of length `nTestNovel` with 2-element tuples. The
+                1st element of each tuple is the image id that was sampled and
+                the 2nd element is its category label (which is in the range
+                [nKbase, nKbase + len(Knovel) - 1]).
+            Exemplars: a list of length len(Knovel) * nExemplars of 2-element
+                tuples. The 1st element of each tuple is the image id that was
+                sampled and the 2nd element is its category label (which is in
+                the ragne [nKbase, nKbase + len(Knovel) - 1]).
+        """
+
+        if len(Knovel) == 0:
+            return [], []
+
+        nKnovel = len(Knovel)
+        Tnovel = []
+        Exemplars = []
+        assert((nTestNovel % nKnovel) == 0)
+        nEvalExamplesPerClass = int(nTestNovel / nKnovel)
+
+        for Knovel_idx in range(len(Knovel)):
+            imd_ids = self.sampleImageIdsFrom(
+                Knovel[Knovel_idx],
+                sample_size=(nEvalExamplesPerClass + nExemplars))
+
+            imds_tnovel = imd_ids[:nEvalExamplesPerClass]
+            imds_ememplars = imd_ids[nEvalExamplesPerClass:]
+
+            Tnovel += [(img_id, nKbase+Knovel_idx) for img_id in imds_tnovel]
+            Exemplars += [(img_id, nKbase+Knovel_idx) for img_id in imds_ememplars]
+        assert(len(Tnovel) == nTestNovel)
+        assert(len(Exemplars) == len(Knovel) * nExemplars)
+
+        # random.shuffle(Exemplars)
+
+        return Tnovel, Exemplars
+
+    def sample_episode(self):
+        """Samples a training episode."""
+        nKnovel = self.nKnovel
+        nKbase = self.nKbase
+        nTestNovel = self.nTestNovel
+        nTestBase = self.nTestBase
+        nExemplars = self.nExemplars
+
+        Kbase, Knovel = self.sample_base_and_novel_categories(nKbase, nKnovel)
+        Tbase = self.sample_test_examples_for_base_categories(Kbase, nTestBase)
+        Tnovel, Exemplars = self.sample_train_and_test_examples_for_novel_categories(
+            Knovel, nTestNovel, nExemplars, nKbase)
+
+        # concatenate the base and novel category examples.
+        Test = Tbase + Tnovel
+        # random.shuffle(Test)
+        Kall = Kbase + Knovel
+
+        return Exemplars, Test, Kall, nKbase
+
+    def createExamplesTensorData(self, examples):
+        """
+        Creates the examples image and label tensor data.
+
+        Args:
+            examples: a list of 2-element tuples, each representing a
+                train or test example. The 1st element of each tuple
+                is the image id of the example and 2nd element is the
+                category label of the example, which is in the range
+                [0, nK - 1], where nK is the total number of categories
+                (both novel and base).
+
+        Returns:
+            images: a tensor of shape [nExamples, Height, Width, 3] with the
+                example images, where nExamples is the number of examples
+                (i.e., nExamples = len(examples)).
+            labels: a tensor of shape [nExamples] with the category label
+                of each example.
+        """
+        images = torch.stack(
+            [self.dataset[img_idx][0] for img_idx, _ in examples], dim=0)
+        labels = torch.LongTensor([label for _, label in examples])
+        return images, labels
+
+    def get_iterator(self, epoch=0):
+        rand_seed = epoch
+        random.seed(rand_seed)
+        np.random.seed(rand_seed)
+        def load_function(iter_idx):
+            Exemplars, Test, Kall, nKbase = self.sample_episode()
+            Xt, Yt = self.createExamplesTensorData(Test)
+            Kall = torch.LongTensor(Kall)
+            if len(Exemplars) > 0:
+                Xe, Ye = self.createExamplesTensorData(Exemplars)
+                return Xe, Ye, Xt, Yt, Kall, nKbase
+            else:
+                return Xt, Yt, Kall, nKbase
+
+        tnt_dataset = tnt.dataset.ListDataset(
+            elem_list=range(self.epoch_size), load=load_function)
+        data_loader = tnt_dataset.parallel(
+            batch_size=self.batch_size,
+            num_workers=(0 if self.is_eval_mode else self.num_workers),
+            shuffle=(False if self.is_eval_mode else True))
+
+        return data_loader
+
+    def __call__(self, epoch=0):
+        return self.get_iterator(epoch)
+
+    def __len__(self):
+        return int(self.epoch_size / self.batch_size)

norm.py

@@ -0,0 +1,32 @@
+import pandas as pd
+import numpy as np
+import cv2
+import statistics
+from tqdm import tqdm
+from glob import glob
+
+def calculate_normalization_parameters(path=None):
+    # data = pd.read_csv(path_to_train_csv)
+    data = glob('NIH/images/*.png')
+    mean = 0
+    std = 0
+    height = []
+    width = []
+    for i in tqdm(data):
+        image = cv2.imread(i)[:, :, ::-1]
+        h, w, _ = image.shape
+        image = image.reshape(-1, 3)
+        mean += np.mean(image, axis=0)
+        std += np.std(image, axis=0)
+        height.append(h)
+        width.append(w)
+    mean = mean / (255 * len(data))
+    std = std / (255 * len(data))
+    print("median height:", statistics.median(height))
+    print("median width:", statistics.median(width))
+    print("mean:", mean)
+    print("std:", std)
+    return mean, std
+
+
+calculate_normalization_parameters()

requirements.txt

@@ -0,0 +1,10 @@
+einops
+timm 
+torchinfo 
+torchsummary 
+torchnet
+wandb
+adabelief_pytorch 
+scikit-plot
+pandas
+h5py

test.py

@@ -0,0 +1,345 @@
+# -*- coding: utf-8 -*-
+import argparse
+
+import torch
+import torch.nn.functional as F
+from torch.utils.data import DataLoader
+
+from torch.autograd import Variable
+
+from tqdm import tqdm
+
+from models.protonet_embedding import ProtoNetEmbedding
+from models.R2D2_embedding import R2D2Embedding
+from models.ResNet12_embedding import resnet12
+
+from models.classification_heads import ClassificationHead
+
+from utils import pprint, set_gpu, Timer, count_accuracy, log
+from sklearn.metrics import confusion_matrix, f1_score, roc_curve, auc
+import scikitplot as skplt
+import matplotlib.pyplot as plt
+
+
+import numpy as np
+import os
+import random
+
+import pickle 
+
+from dataloader.chest import label_dict
+
+
+import pandas as pd
+
+def multiclass_roc(y_test, y_score,n_classes = 3):
+    
+
+    # structures
+    fpr = dict()
+    tpr = dict()
+    roc_auc = dict()
+
+    # calculate dummies once
+    y_test_dummies = pd.get_dummies(y_test, drop_first=False).values
+    for i in range(n_classes):
+        fpr[i], tpr[i], _ = roc_curve(y_test_dummies[:, i], y_score[:, i])
+        roc_auc[i] = auc(fpr[i], tpr[i])
+
+    return fpr,tpr,roc_auc
+
+# os.environ['CUDA_VISIBLE_DEVICES'] = "0"
+
+def seed_everything(seed: int):
+    random.seed(seed)
+    os.environ["PYTHONHASHSEED"] = str(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cudnn.benchmark = True
+
+def euclidean_dist(x, y):
+
+    # x: N x D
+    # y: M x D
+    n = x.size(0)
+    m = y.size(0)
+    d = x.size(1)
+
+    assert d == y.size(1)
+
+    x = x.unsqueeze(1).expand(n, m, d)
+    y = y.unsqueeze(0).expand(n, m, d)
+
+    return torch.pow(x - y, 2).sum(2)
+
+def flip(x, dim):
+    xsize = x.size()
+    dim = x.dim() + dim if dim < 0 else dim
+    x = x.view(-1, *xsize[dim:])
+    x = x.view(x.size(0), x.size(1), -1)[:, getattr(torch.arange(x.size(1)-1,
+                      -1, -1), ('cpu','cuda')[x.is_cuda])().long(), :]
+    return x.view(xsize)
+
+
+def get_model(options):
+    # Choose the embedding network
+    if options.network == 'ProtoNet':
+        network = ProtoNetEmbedding().cuda()
+    elif options.network == 'R2D2':
+        network = R2D2Embedding().cuda()
+    elif options.network == 'ResNet':
+        if options.dataset == 'miniImageNet' or options.dataset == 'tieredImageNet':
+            network = resnet12(avg_pool=False, drop_rate=0.1, dropblock_size=5,num_layer=options.num_layer).cuda()
+            network = torch.nn.DataParallel(network)
+        else:
+            network = resnet12(avg_pool=False, drop_rate=0.1, dropblock_size=2,num_layer=options.num_layer).cuda()
+    else:
+        print ("Cannot recognize the network type")
+        assert(False)
+
+    # Choose the classification head
+    if opt.head == 'ProtoNet':
+        cls_head = ClassificationHead(base_learner='ProtoNet').cuda()
+    elif options.head == 'SubspaceTrans':
+        cls_head = ClassificationHead(base_learner='SubspaceTrans').cuda()
+    elif options.head == 'Subspace':
+        cls_head = ClassificationHead(base_learner='Subspace').cuda()
+    elif options.head == 'SubspaceFast':
+        cls_head = ClassificationHead(base_learner='SubspaceFast').cuda()
+    elif opt.head == 'Ridge':
+        cls_head = ClassificationHead(base_learner='Ridge').cuda()
+    elif opt.head == 'R2D2':
+        cls_head = ClassificationHead(base_learner='R2D2').cuda()
+    elif opt.head == 'SVM':
+        cls_head = ClassificationHead(base_learner='SVM-CS').cuda()
+    else:
+        print ("Cannot recognize the classification head type")
+        assert(False)
+
+    return (network, cls_head)
+
+def get_dataset(options):
+    # Choose the embedding network
+    if options.dataset == 'miniImageNet':
+        from dataloader.mini_imagenet import MiniImageNet, FewShotDataloader
+        dataset_test = MiniImageNet(phase='test')
+        data_loader = FewShotDataloader
+    elif options.dataset == 'tieredImageNet':
+        from dataloader.tiered_imagenet import tieredImageNet, FewShotDataloader
+        dataset_test = tieredImageNet(phase='test')
+        data_loader = FewShotDataloader
+    elif options.dataset == 'CIFAR_FS':
+        from dataloader.CIFAR_FS import CIFAR_FS, FewShotDataloader
+        dataset_test = CIFAR_FS(phase='test')
+        data_loader = FewShotDataloader
+    elif options.dataset == 'FC100':
+        from dataloader.FC100 import FC100, FewShotDataloader
+        dataset_test = FC100(phase='test')
+        data_loader = FewShotDataloader
+    elif options.dataset == 'Chest':
+        from dataloader.chest import Chest, FewShotDataloader
+        dataset_test = Chest(phase='test')
+        data_loader = FewShotDataloader
+    else:
+        print ("Cannot recognize the dataset type")
+        assert(False)
+
+    return (dataset_test, data_loader)
+
+#
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+
+    #Changes
+    parser.add_argument('--gpu', default='3')
+    #Changes    
+    parser.add_argument('--load', 
+                        default='experiments/group2_subspace30_CE_train/best_model.pth', ## your best model
+                        help='path of the checkpoint file')
+    #Changes
+    parser.add_argument('--num_layer', type=int, default=30,
+                            help='num of layer')
+    
+    parser.add_argument('--episode', type=int, default=1000,
+                            help='number of episodes to test')
+    parser.add_argument('--way', type=int, default=3,
+                            help='number of classes in one test episode')
+    parser.add_argument('--shot', type=int, default=5,
+                            help='number of support examples per training class')
+    parser.add_argument('--query', type=int, default=5,
+                            help='number of query examples per training class')
+    parser.add_argument('--network', type=str, default='ResNet',
+                            help='choose which embedding network to use. ProtoNet, R2D2, ResNet')
+    parser.add_argument('--head', type=str, default='Subspace',
+                            help='choose which embedding network to use. ProtoNet, Ridge, R2D2, SVM')
+    parser.add_argument('--dataset', type=str, default='Chest',
+                            help='choose which classification head to use. miniImageNet, tieredImageNet, CIFAR_FS, FC100')
+
+
+    opt = parser.parse_args()
+
+    seed_everything(42)
+
+    (dataset_test, data_loader) = get_dataset(opt)
+
+    set_gpu(opt.gpu)
+
+    # Define the models
+    (embedding_net, cls_head) = get_model(opt)
+
+    # Load saved model checkpoints
+    saved_models = torch.load(opt.load)
+    embedding_net.load_state_dict(saved_models['embedding'])
+    embedding_net.eval()
+    cls_head.load_state_dict(saved_models['head'])
+    cls_head.eval()
+
+
+    aug=False
+
+    label_dict_inv =  {v:k for k,v in label_dict.items()}
+
+    test_accuracies = []
+    per_class_accuracies = []
+    y_pred_list = []
+    y_list = []
+    dloader_test = data_loader(
+        dataset=dataset_test,
+        nKnovel=opt.way,
+        nKbase=0,
+        nExemplars=opt.shot,  # num training examples per novel category
+        nTestNovel=opt.query * opt.way,  # num test examples for all the novel categories
+        nTestBase=0,  # num test examples for all the base categories
+        batch_size=1,
+        num_workers=1,
+        epoch_size=opt.episode,  # num of batches per epoch
+    )
+
+    #print("epp:  ", epp)
+
+    with torch.no_grad():
+        for i, batch in enumerate(tqdm(dloader_test()), 1):
+            data_support, labels_support, data_query, labels_query, _, _ = [x.cuda() for x in batch]
+
+            n_support = opt.way * opt.shot
+            n_query = opt.way * opt.query
+
+            if opt.shot == 1 and aug:
+                flipped_data_support = flip(data_support, 3)
+                data_support = torch.cat((data_support, flipped_data_support), dim=0)
+                labels_support = torch.cat((labels_support, labels_support), dim=0)
+
+            list_emb_support = embedding_net(data_support.reshape([-1] + list(data_support.shape[-3:])))
+            list_emb_query = embedding_net(data_query.reshape([-1] + list(data_query.shape[-3:])))
+
+            logits = torch.zeros(n_query, opt.way).cuda()
+
+            for emb_support, emb_query in zip(list_emb_support, list_emb_query):
+
+
+                emb_support = emb_support.view(1, opt.way, opt.shot, -1).mean(2)
+
+                emb_query = emb_query.reshape(1, n_query, -1)
+
+                dists = euclidean_dist(emb_query[0], emb_support[0])
+
+
+                logits += F.softmax(-dists, dim=1).view(1 * opt.way * opt.query, -1)
+
+
+
+            logits /= opt.num_layer
+
+            logits = logits.reshape(-1, opt.way)
+            labels_query = labels_query.reshape(-1)
+
+
+            acc,pca = count_accuracy(logits, labels_query)
+            test_accuracies.append(acc.item())
+            per_class_accuracies.append(pca)
+
+            y_pred_list.append(logits.detach().cpu().numpy())
+            y_list.append(labels_query.detach().cpu().numpy())
+
+            avg = np.mean(np.array(test_accuracies))
+            std = np.std(np.array(test_accuracies))
+            ci95 = 1.96 * std / np.sqrt(i + 1)
+
+            if i % 10 == 0:
+                
+                # print(logits.detach().cpu().numpy())
+                # print(torch.argmax(logits, dim=1).view(-1))
+                # print(labels_query.detach().cpu().numpy())
+
+                pca = np.array(per_class_accuracies).mean(0)
+                pcs = np.array(per_class_accuracies).std(0)
+
+                print('Episode [{}/{}]:\t\t\tAccuracy: {:.2f} ± {:.2f} ({:.2f}) % ({:.2f} %)'\
+                    .format(i, opt.episode, avg, ci95,std, acc))
+                print(f'{label_dict_inv[9]}: {pca[0]:.2f} ± {pcs[0]:.2f} % | {label_dict_inv[10]}: {pca[1]:.2f} ± {pcs[1]:.2f} % | {label_dict_inv[11]}: {pca[2]:.2f} ± {pcs[2]:.2f}%')
+
+        
+
+    pca = np.array(per_class_accuracies).mean(0)
+    pcs = np.array(per_class_accuracies).std(0)
+
+    print("Mean")
+    print(pca)
+    print('Standard Deviation')
+    print(pcs)
+
+
+    y_pred_proba = np.array(
+        y_pred_list).reshape(-1, 3)
+
+    y_pred = np.argmax(y_pred_proba, axis=1)
+
+    y_true = np.array(y_list).reshape(-1)
+
+    f1 = f1_score(y_true, y_pred, average=None)
+
+    print('F1 Score')
+    print(f1)
+
+    fpr,tpr, auc = multiclass_roc(y_true,y_pred_proba)
+    save_tuple = (fpr,tpr,auc)
+
+    print(auc)
+
+    # Plots
+
+    #Changes
+    # with open('plot/group5_subspace25.pickle', 'wb') as f:
+    #     pickle.dump(save_tuple, f)
+
+    #Changes
+    class_dict = {'Fibrosis': 0, 'Hernia': 1, 'Pneumonia': 2}
+    # class_dict = {'Mass': 0, 'Nodule': 1, 'Pleural_Thickening': 2}
+    # class_dict = {'Cardiomegaly': 0, 'Edema': 1, 'Emphysema': 2}
+    # class_dict = {'Consolidation': 0, 'Effusion': 1, 'Pneumothorax': 2}
+    # class_dict = {'Atelectasis': 0, 'Infiltration': 1, 'No Finding': 2}
+
+    class_dict_inv = {v: k for k, v in class_dict.items()}
+
+    y_true = np.array([class_dict_inv[i]
+                      for i in np.array(y_list).reshape(-1)])
+
+    # print(np.array(y_pred_list).reshape(-1, 3).shape)
+    # print(np.array(y_list).reshape(-1).shape)
+    # print(y_list)
+    # print(np.array(y_pred_list).reshape(-1, 3))
+
+
+    # skplt.metrics.plot_roc(y_true, y_pred_proba,plot_micro=False, plot_macro=False)
+
+    #Changes
+    # plt.savefig('plot/group5_subspace25.png', dpi=1000)
+    # plt.show()
+
+
+
+
+# python test_ortho_bcs.py --gpu 2 --load experiments/chest_exp1/best_model.pth --way 3 --dataset Chest
+# python test_ortho_bcs.py --gpu 2 --load experiments/chest_exp1/best_model.pth --way 3 --dataset Chest

train.py

@@ -0,0 +1,458 @@
+# -*- coding: utf-8 -*-
+import timm
+import os
+import sys
+import argparse
+import random
+import numpy as np
+from tqdm import tqdm
+import torch
+import torch.nn.functional as F
+from torch import linalg as LA
+from models.classification_heads import ClassificationHead
+from models.R2D2_embedding import R2D2Embedding
+from models.protonet_embedding import ProtoNetEmbedding
+from models.ResNet12_embedding import resnet12
+import torch.nn as nn
+from utils import set_gpu, Timer, count_accuracy, check_dir, log
+import warnings
+import wandb
+from itertools import combinations
+
+from torchsummary import summary
+warnings.filterwarnings("ignore")
+
+
+def one_hot(indices, depth):
+    """
+    Returns a one-hot tensor.
+    This is a PyTorch equivalent of Tensorflow's tf.one_hot.
+
+    Parameters:
+      indices:  a (n_batch, m) Tensor or (m) Tensor.
+      depth: a scalar. Represents the depth of the one hot dimension.
+    Returns: a (n_batch, m, depth) Tensor or (m, depth) Tensor.
+    """
+
+    encoded_indicies = torch.zeros(indices.size() + torch.Size([depth])).cuda()
+    index = indices.view(indices.size()+torch.Size([1]))
+    encoded_indicies = encoded_indicies.scatter_(1, index, 1)
+
+    return encoded_indicies
+
+def seed_everything(seed: int):
+    random.seed(seed)
+    os.environ["PYTHONHASHSEED"] = str(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cudnn.benchmark = True
+
+
+def euclidean_dist(x, y):
+
+    # x: N x D
+    # y: M x D
+    n = x.size(0)
+    m = y.size(0)
+    d = x.size(1)
+
+    assert d == y.size(1)
+
+    x = x.unsqueeze(1).expand(n, m, d)
+    y = y.unsqueeze(0).expand(n, m, d)
+
+
+    return torch.pow(x - y, 2).sum(2)
+
+def cosine_dist(x, y):
+
+    # x: N x D
+    # y: M x D
+    n = x.size(0)
+    m = y.size(0)
+    d = x.size(1)
+
+    assert d == y.size(1)
+
+    x = x.unsqueeze(1).expand(n, m, d)
+    y = y.unsqueeze(0).expand(n, m, d)
+
+
+
+    cos = nn.CosineSimilarity(dim=2, eps=1e-6)
+    out = 1 - cos(x,y)
+
+
+    return out
+
+
+def get_model(options):
+    # Choose the embedding network
+    if options.network == 'ProtoNet':
+        network = ProtoNetEmbedding().cuda()
+    elif options.network == 'R2D2':
+        network = R2D2Embedding().cuda()
+    elif options.network == 'ResNet':
+        if options.dataset == 'miniImageNet' or options.dataset == 'tieredImageNet':
+            network = resnet12(avg_pool=False, drop_rate=0.1,
+                               dropblock_size=5,num_layer=options.num_layer).cuda()
+            network = torch.nn.DataParallel(network)  # , device_ids=[1, 2])
+        else:
+            network = resnet12(avg_pool=False, drop_rate=0.1,
+                               dropblock_size=2,num_layer=options.num_layer).cuda()
+    else:
+        print("Cannot recognize the network type")
+        assert(False)
+
+    # Choose the classification head
+    if options.head == 'Subspace':
+        cls_head = ClassificationHead(base_learner='Subspace').cuda()
+    elif options.head == 'ProtoNet':
+        cls_head = ClassificationHead(base_learner='ProtoNet').cuda()
+    elif options.head == 'Ridge':
+        cls_head = ClassificationHead(base_learner='Ridge').cuda()
+    elif options.head == 'R2D2':
+        cls_head = ClassificationHead(base_learner='R2D2').cuda()
+    elif options.head == 'SVM':
+        cls_head = ClassificationHead(base_learner='SVM-CS').cuda()
+    else:
+        print("Cannot recognize the dataset type")
+        assert(False)
+
+    return (network, cls_head)
+
+def get_dataset(options):
+    # Choose the embedding network
+    if options.dataset == 'miniImageNet':
+        from dataloader.mini_imagenet import MiniImageNet, FewShotDataloader
+        # change it to train only, this is including the validation set
+        dataset_train = MiniImageNet(phase='trainval')
+        dataset_val = MiniImageNet(phase='test')
+        data_loader = FewShotDataloader
+    elif options.dataset == 'tieredImageNet':
+        from dataloader.tiered_imagenet import tieredImageNet, FewShotDataloader
+        dataset_train = tieredImageNet(phase='train')
+        dataset_val = tieredImageNet(phase='test')
+        data_loader = FewShotDataloader
+    elif options.dataset == 'CIFAR_FS':
+        from dataloader.CIFAR_FS import CIFAR_FS, FewShotDataloader
+        dataset_train = CIFAR_FS(phase='train')
+        dataset_val = CIFAR_FS(phase='test')
+        data_loader = FewShotDataloader
+    elif options.dataset == 'Chest':
+        from dataloader.chest import Chest, FewShotDataloader
+        dataset_train = Chest(phase='train')
+        dataset_val = Chest(phase='val')
+        data_loader = FewShotDataloader
+    else:
+        print("Cannot recognize the dataset type")
+        assert(False)
+
+    return (dataset_train, dataset_val, data_loader)
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--num-epoch', type=int, default=80,
+                        help='number of training epochs')
+    parser.add_argument('--save-epoch', type=int, default=5,
+                        help='frequency of model saving')
+    parser.add_argument('--train-shot', type=int, default=5,
+                        help='number of support examples per training class')
+    parser.add_argument('--val-shot', type=int, default=5,
+                        help='number of support examples per validation class')
+    parser.add_argument('--train-query', type=int, default=5,
+                        help='number of query examples per training class')
+    parser.add_argument('--val-episode', type=int, default=600,
+                        help='number of episodes per validation')
+    parser.add_argument('--val-query', type=int, default=5,
+                        help='number of query examples per validation class')
+    parser.add_argument('--train-way', type=int, default=3,
+                        help='number of classes in one training episode')
+    parser.add_argument('--test-way', type=int, default=3,
+                        help='number of classes in one test (or validation) episode')
+    parser.add_argument('--save-path', default='experiments')
+
+    parser.add_argument('--wandbexperiment', default="group5_subspace30",type=str)
+    parser.add_argument('--gpu', default='0')  # using 4 gpus
+    parser.add_argument('--num_layer', type=int, default=30,
+                        help='number of linear layer')
+
+    # parser.add_argument('--gpu', default='0,1,2,3')  # using 4 gpus
+    parser.add_argument('--network', type=str, default='ResNet',
+                        help='choose which embedding network to use. ResNet')
+    parser.add_argument('--head', type=str, default='Subspace',
+                        help='choose which classification head to use. Subspace, ProtoNet, R2D2, SVM')
+    parser.add_argument('--dataset', type=str, default='Chest',
+                        help='choose which classification head to use. miniImageNet, tieredImageNet, CIFAR_FS, FC100')
+    parser.add_argument('--episodes-per-batch', type=int, default=1,
+                        help='number of episodes per batch')
+    parser.add_argument('--eps', type=float, default=0.0,
+                        help='epsilon of label smoothing')
+    parser.add_argument('--wandb', action="store_true")
+    parser.add_argument("--wandbkey", type=str,
+                        default='db1158429a436f94565ac9eadecc6afe9e5a0b8f',
+                        help='Wandb project key')
+
+
+# python train_my.py --gpu 2 --dataset Chest --num_layer 5
+
+
+    opt = parser.parse_args()
+    seed_everything(42)
+    print(opt)
+    opt.save_path = os.path.join(opt.save_path,opt.wandbexperiment)
+
+
+    if opt.wandb:
+        os.system('wandb login {}'.format(opt.wandbkey))
+        wandb.init(name=opt.wandbexperiment,
+                   project='chest-few-shot-final')
+        wandb.config.update(opt)
+
+    (dataset_train, dataset_val, data_loader) = get_dataset(opt)
+
+    # Dataloader of Gidaris & Komodakis (CVPR 2018)
+    dloader_train = data_loader(
+        dataset=dataset_train,
+        nKnovel=opt.train_way,
+        nKbase=0,
+        nExemplars=opt.train_shot,  # num training examples per novel category
+        # num test examples for all the novel categories
+        nTestNovel=opt.train_way * opt.train_query,
+        nTestBase=0,  # num test examples for all the base categories
+        batch_size=opt.episodes_per_batch,
+        num_workers=15,
+        epoch_size=opt.episodes_per_batch * 1000,  # num of batches per epoch
+    )
+
+    dloader_val = data_loader(
+        dataset=dataset_val,
+        nKnovel=opt.test_way,
+        nKbase=0,
+        nExemplars=opt.val_shot,  # num training examples per novel category
+        # num test examples for all the novel categories
+        nTestNovel=opt.val_query * opt.test_way,
+        nTestBase=0,  # num test examples for all the base categories
+        batch_size=1,
+        num_workers=15,
+        epoch_size=1 * opt.val_episode,  # num of batches per epoch
+    )
+
+    set_gpu(opt.gpu)
+    check_dir('./experiments/')
+    check_dir(opt.save_path)
+
+    log_file_path = os.path.join(opt.save_path, "train_log.txt")
+    log(log_file_path, str(vars(opt)))
+
+    (embedding_net, cls_head) = get_model(opt)
+
+    optimizer = torch.optim.SGD(embedding_net.parameters(),lr=3e-3)
+
+
+    def lambda_epoch(e): return 1.0 if e < 12 else (
+        0.025 if e < 30 else 0.0032 if e < 45 else (0.0014 if e < 57 else (0.00052)))
+
+    ## tieredimagenet###
+    # lambda_epoch = lambda e: 1.0 if e < 20 else (
+    #     0.012 if e < 45 else 0.0052 if e < 59 else (0.00054 if e < 68 else (0.00012)))
+
+    lr_scheduler = torch.optim.lr_scheduler.LambdaLR(
+        optimizer, lr_lambda=lambda_epoch, last_epoch=-1)
+
+    max_val_acc = 0.0
+
+    timer = Timer()
+    x_entropy = torch.nn.CrossEntropyLoss()
+
+
+    index = list(combinations([i for i in range(opt.num_layer)], 2))
+
+    for epoch in range(1, opt.num_epoch + 1):
+
+
+        for param_group in optimizer.param_groups:
+            epoch_learning_rate = param_group['lr']
+
+        log(log_file_path, 'Train Epoch: {}\tLearning Rate: {:.4f}'.format(
+            epoch, epoch_learning_rate))
+
+        _, _ = [x.train() for x in (embedding_net, cls_head)]
+
+        train_accuracies = []
+        train_losses = []
+
+        train_n_support = opt.train_way * opt.train_shot
+        train_n_query = opt.train_way * opt.train_query
+
+
+
+
+        for i, batch in enumerate(tqdm(dloader_train(epoch)), 1):
+
+            data_support, labels_support, data_query, labels_query, _, _ = [
+                x.cuda() for x in batch]
+
+            list_emb_query = embedding_net(data_query.view(
+                [-1] + list(data_query.shape[-3:])))  # [100, 2560]
+            list_emb_support  = embedding_net(data_support.view(
+                [-1] + list(data_support.shape[-3:])))  # [100, 3, 32, 32] -> [100, 2560]
+
+
+            loss_weights = 0.
+            for ind in index:
+
+                loss_weights += torch.abs(F.cosine_similarity(getattr(embedding_net,f'linear{ind[0]}_1').weight.view(-1),getattr(embedding_net,f'linear{ind[1]}_1').weight.view(-1),dim=0))
+
+
+            log_p_y = torch.zeros(
+                opt.episodes_per_batch * opt.train_way * opt.train_query, opt.train_way).cuda()
+
+            for emb_support,emb_query in zip(list_emb_support, list_emb_query):
+                # emb_support = emb_support.view(
+                #     opt.episodes_per_batch, train_n_support, -1)  # [4, 25, 2560]
+                if opt.train_shot == 1:
+                    emb_support = emb_support.view(
+                        opt.episodes_per_batch, opt.train_way, -1)  # [4,5,5,2560] --> [4, 5, 20]
+                else:
+                    emb_support = emb_support.view(
+                        opt.episodes_per_batch, opt.train_way, opt.train_shot, -1).mean(2)  # [4,5,5,2560] --> [4, 5, 20]
+
+                emb_query = emb_query.view(
+                    opt.episodes_per_batch, train_n_query, -1)  # [4, 25, 2560]
+
+
+                dists = torch.stack(
+                    [euclidean_dist(emb_query[i], emb_support[i]) for i in range(opt.episodes_per_batch)])  # [4,25,5]
+
+
+
+                log_p_y += F.softmax(-dists,
+                                     dim=2).view(opt.episodes_per_batch* opt.train_way* opt.train_query, -1)  # [100,5]
+
+
+            log_p_y /= opt.num_layer
+
+
+            smoothed_one_hot = one_hot(
+                labels_query.view(-1), opt.train_way)  # [100,5]
+
+            loss = x_entropy(
+                log_p_y.view(-1, opt.train_way), labels_query.view(-1))
+
+
+            acc, _ = count_accuracy(
+                log_p_y.view(-1, opt.train_way), labels_query.view(-1))
+
+            train_accuracies.append(acc.item())
+            train_losses.append(loss.item())
+
+            if (i % 100 == 0):
+                train_acc_avg = np.mean(np.array(train_accuracies))
+                log(log_file_path, 'Train Epoch: {}\tBatch: [{}/{}]\tLoss: {:.4f}\tAccuracy: {:.2f} % ({:.2f} %)'.format(
+                    epoch, i, len(dloader_train), loss.item(), train_acc_avg, acc))
+                if opt.wandb:
+
+                    wandb.log({'Epoch': epoch,
+                           'lr': optimizer.param_groups[0]['lr'],"Loss":loss.item(),"Avg Accuracy":train_acc_avg,'Accuracy':acc,
+                           'cosine loss':loss_weights})
+
+
+            optimizer.zero_grad()
+
+            loss += loss_weights
+            loss.backward()
+
+            optimizer.step()
+
+        # Evaluate on the validation split
+        _, _ = [x.eval() for x in (embedding_net, cls_head)]
+
+        val_accuracies = []
+        val_losses = []
+
+
+        with torch.no_grad():
+
+            for i, batch in enumerate(tqdm(dloader_val(epoch)), 1):
+                data_support, labels_support, data_query, labels_query, _, _ = [
+                    x.cuda() for x in batch]
+
+                test_n_support = opt.test_way * opt.val_shot
+                test_n_query = opt.test_way * opt.val_query
+
+
+                list_emb_support = embedding_net(data_support.view(
+                    [-1] + list(data_support.shape[-3:])))
+                list_emb_query = embedding_net(data_query.view(
+                    [-1] + list(data_query.shape[-3:])))
+
+
+                logit_query = torch.zeros(test_n_query, opt.test_way).cuda()
+
+                for emb_support, emb_query in zip(list_emb_support, list_emb_query):
+
+                    # print(emb_support.size())
+                    emb_support = emb_support.view(1, test_n_support, -1)
+                    # print(emb_support.size())
+
+                    emb_support = emb_support.view(
+                        1, opt.train_way, opt.train_shot, -1).mean(2)  # [4, 5, 20]
+
+                    emb_query = emb_query.view(1, test_n_query, -1)
+
+                    # print(emb_support.size(),emb_query.size())
+
+                    dists = torch.stack(
+                        [euclidean_dist(emb_query[i], emb_support[i]) for i in range(emb_query.size(0))])
+
+                    logit_query += F.softmax(-dists, dim=2).view(1 *
+                                                                opt.test_way * opt.val_query, -1)  # []
+
+                logit_query /= opt.num_layer
+
+
+                loss = x_entropy(
+                    logit_query.view(-1, opt.test_way), labels_query.view(-1))
+                acc, _ = count_accuracy(
+                    logit_query.view(-1, opt.test_way), labels_query.view(-1))
+
+                val_accuracies.append(acc.item())
+                val_losses.append(loss.item())
+
+        val_acc_avg = np.mean(np.array(val_accuracies))
+        val_acc_ci95 = 1.96 * \
+            np.std(np.array(val_accuracies)) / np.sqrt(opt.val_episode)
+
+        val_loss_avg = np.mean(np.array(val_losses))
+
+        if val_acc_avg > max_val_acc:
+            max_val_acc = val_acc_avg
+            torch.save({'embedding': embedding_net.state_dict(), 'head': cls_head.state_dict()},
+                       os.path.join(opt.save_path, 'best_model.pth'))
+
+
+
+            log(log_file_path, 'Validation Epoch: {}\t\t\tLoss: {:.4f}\tAccuracy: {:.2f} ± {:.2f} % (Best)'
+                .format(epoch, val_loss_avg, val_acc_avg, val_acc_ci95))
+        else:
+            log(log_file_path, 'Validation Epoch: {}\t\t\tLoss: {:.4f}\tAccuracy: {:.2f} ± {:.2f} %'
+                .format(epoch, val_loss_avg, val_acc_avg, val_acc_ci95))
+
+        if opt.wandb:
+            wandb.log({"Validation Loss":val_loss_avg,"Val Avg Accuracy":val_acc_avg})
+
+        torch.save({'embedding': embedding_net.state_dict(
+        ), 'head': cls_head.state_dict()}, os.path.join(opt.save_path, 'last_epoch.pth'))
+
+        if epoch % opt.save_epoch == 0:
+            torch.save({'embedding': embedding_net.state_dict(), 'head': cls_head.state_dict(
+            )}, os.path.join(opt.save_path, 'epoch_{}.pth'.format(epoch)))
+
+        log(log_file_path, 'Elapsed Time: {}/{}\n'.format(timer.measure(),
+                                                          timer.measure(epoch / float(opt.num_epoch))))
+
+    # lr_scheduler.step()

utils.py

@@ -0,0 +1,56 @@
+import os
+import time
+import pprint
+import torch
+from sklearn.metrics import confusion_matrix
+
+def set_gpu(x):
+    os.environ['CUDA_VISIBLE_DEVICES'] = x
+    print('using gpu:', x)
+
+def check_dir(path):
+    '''
+    Create directory if it does not exist.
+        path:           Path of directory.
+    '''
+    if not os.path.exists(path):
+        os.mkdir(path)
+
+def count_accuracy(logits, label):
+    pred = torch.argmax(logits, dim=1).view(-1)
+    label = label.view(-1)
+
+    acc = [0 for c in range(3)]  
+    for c in range(3):
+        acc[c] = (pred.eq(label) * label.eq(c)).float() / max((label.eq(c)).sum(), 1)
+
+
+    matrix = confusion_matrix(label.cpu().detach().numpy(), pred.cpu().detach().numpy())
+    pca = matrix.diagonal()/matrix.sum(axis=1)
+
+    accuracy = 100 * pred.eq(label).float().mean()
+    return accuracy, pca * 100
+    
+class Timer():
+    def __init__(self):
+        self.o = time.time()
+
+    def measure(self, p=1):
+        x = (time.time() - self.o) / float(p)
+        x = int(x)
+        if x >= 3600:
+            return '{:.1f}h'.format(x / 3600)
+        if x >= 60:
+            return '{}m'.format(round(x / 60))
+        return '{}s'.format(x)
+
+def log(log_file_path, string):
+    '''
+    Write one line of log into screen and file.
+        log_file_path: Path of log file.
+        string:        String to write in log file.
+    '''
+    with open(log_file_path, 'a+') as f:
+        f.write(string + '\n')
+        f.flush()
+    print(string)