# -----------------
# DATASET ROOTS
# -----------------
cifar_10_root = '/disk/scratch_fast/bingchen/gcd/cifar10'
cifar_100_root = '/disk/scratch_fast/bingchen/gcd/cifar100'
cub_root = '/disk/scratch_fast/bingchen/ssb/cub'
nabirds_root = '/disk/scratch_fast/bingchen/nabirds'
inat21_root = '/disk/scratch_fast/bingchen/inat21'
inat18_root = '/disk/scratch_fast/bingchen/inat18'
inat19_root = '/disk/scratch_fast/bingchen/inat19'
aircraft_root = '/disk/scratch_fast/bingchen/ssb/aircraft/fgvc-aircraft-2013b'
car_root = '/disk/scratch_fast/bingchen/ssb/cars'
herbarium_dataroot = '/disk/scratch_fast/bingchen/herb19'
imagenet_root = '/disk/scratch_fast/bingchen/imagenet1k/ImageNet1k'
sun_root = '/disk/scratch_big/bingchen/dataset/alternative_grouping/SUN'

# OSR Split dir
osr_split_dir = 'data/ssb_splits'

# -----------------
# OTHER PATHS
# -----------------
exp_root = 'dev_outputs' # All logs and checkpoints will be saved here
config.py
DTennant/IGCD
from torchvision import transforms

import torch

def get_transform(transform_type='imagenet', image_size=32, args=None):

    if transform_type == 'imagenet':

        mean = (0.485, 0.456, 0.406)
        std = (0.229, 0.224, 0.225)
        interpolation = args.interpolation
        crop_pct = args.crop_pct

        train_transform = transforms.Compose([
            transforms.Resize(int(image_size / crop_pct), interpolation),
            transforms.RandomCrop(image_size),
            transforms.RandomHorizontalFlip(p=0.5),
            transforms.ColorJitter(),
            transforms.ToTensor(),
            transforms.Normalize(
                mean=torch.tensor(mean),
                std=torch.tensor(std))
        ])

        test_transform = transforms.Compose([
            transforms.Resize(int(image_size / crop_pct), interpolation),
            transforms.CenterCrop(image_size),
            transforms.ToTensor(),
            transforms.Normalize(
                mean=torch.tensor(mean),
                std=torch.tensor(std))
        ])
    
    else:

        raise NotImplementedError

    return (train_transform, test_transform)
data/augmentations/__init__.py
DTennant/IGCD
from torchvision.datasets import CIFAR10, CIFAR100
from copy import deepcopy
import numpy as np
import torch

from data.data_utils import subsample_instances
from config import cifar_10_root, cifar_100_root


class CustomCIFAR10(CIFAR10):

    def __init__(self, *args, **kwargs):

        super(CustomCIFAR10, self).__init__(*args, **kwargs)

        self.uq_idxs = np.array(range(len(self)))

    def __getitem__(self, item):

        img, label = super().__getitem__(item)
        uq_idx = self.uq_idxs[item]

        return img, label, uq_idx

    def __len__(self):
        return len(self.targets)


class CustomCIFAR100(CIFAR100):

    def __init__(self, *args, **kwargs):
        super(CustomCIFAR100, self).__init__(*args, **kwargs)

        self.uq_idxs = np.array(range(len(self)))

    def __getitem__(self, item):
        img, label = super().__getitem__(item)
        uq_idx = self.uq_idxs[item]

        return img, label, uq_idx

    def __len__(self):
        return len(self.targets)

class PartialCIFAR100(CIFAR100):
    """`CIFAR100 <https://www.cs.toronto.edu/~kriz/cifar.html>`_ Dataset.
    Only returns the data within the target list and select examples based on the partition.

    """
    def __init__(self, root, split='train', transform=None, target_transform=None, download=False,
                 target_list=range(50), partition=(0, 1.0)):

        super(PartialCIFAR100, self).__init__(root, split == 'train', transform, target_transform, download)

        ind = [
            i for i in range(len(self.targets))
            if self.targets[i] in target_list
        ]
        self.target_list = target_list
        self.data = self.data[ind]
        self.targets = np.array(self.targets)
        self.targets = self.targets[ind].tolist()

        self.targets = torch.tensor(self.targets).long()

        data_of_each_class = []
        label_of_each_class = []
        for label in target_list:
            data_of_cur_class = self.data[self.targets == label]
            label_of_cur_class = self.targets[self.targets == label]
            lower, upper = partition
            lower_idx = int(lower * len(data_of_cur_class)) if lower is not None else 0
            upper_idx = int(upper * len(data_of_cur_class)) if upper is not None else len(data_of_cur_class)
            data_of_each_class.append(data_of_cur_class[lower_idx: upper_idx])
            label_of_each_class.append(label_of_cur_class[lower_idx: upper_idx])
        self.data = np.concatenate(data_of_each_class, axis=0)
        self.targets = torch.cat(label_of_each_class, dim=0)
        
        self.uq_idxs = np.array(range(len(self)))
        
    def concat(self, other_dataset):
        self.data = np.concatenate([self.data, other_dataset.data], axis=0)
        self.targets = torch.cat([self.targets, other_dataset.targets], dim=0)
        
        self.uq_idxs = np.array(range(len(self)))

    def __getitem__(self, item):
        img, label = super().__getitem__(item)
        uq_idx = self.uq_idxs[item]

        return img, label, uq_idx

    def __len__(self):
        return len(self.targets)

def subsample_dataset(dataset, idxs):

    # Allow for setting in which all empty set of indices is passed

    if len(idxs) > 0:

        dataset.data = dataset.data[idxs]
        dataset.targets = np.array(dataset.targets)[idxs].tolist()
        dataset.uq_idxs = dataset.uq_idxs[idxs]

        return dataset

    else:

        return None


def subsample_dataset_lt(dataset, args=None):
    cls_num = len(np.unique(np.array(dataset.targets)))
    def get_img_num_per_cls(dataset, imb_type='exp', imb_factor=0.02):
        img_max = len(dataset) / cls_num
        img_num_per_cls = []
        if imb_type == 'exp':
            for cls_idx in range(cls_num):
                num = img_max * (imb_factor**(cls_idx / (cls_num - 1.0)))
                img_num_per_cls.append(int(num))
        elif imb_type == 'step':
            for cls_idx in range(cls_num // 2):
                img_num_per_cls.append(int(img_max))
            for cls_idx in range(cls_num // 2):
                img_num_per_cls.append(int(img_max * imb_factor))
        else:
            img_num_per_cls.extend([int(img_max)] * cls_num)
        return img_num_per_cls
    
    def gen_imbalanced_data(dataset, img_num_per_cls):
        new_idxs = []
        targets_np = np.array(dataset.targets, dtype=np.int64)
        classes = np.unique(targets_np)
        dataset.num_per_cls_dict = dict()
        for the_class, the_img_num in zip(classes, img_num_per_cls):
            dataset.num_per_cls_dict[the_class] = the_img_num
            idx = np.where(targets_np == the_class)[0]
            np.random.shuffle(idx)
            selec_idx = idx[:the_img_num]
            new_idxs.append(selec_idx)
        new_idxs = np.concatenate(new_idxs, axis=0)
        return new_idxs

    img_num_per_cls = get_img_num_per_cls(dataset, imb_type='exp', imb_factor=args.imbalance_factor)
    new_idxs = gen_imbalanced_data(dataset, img_num_per_cls)
    dataset = subsample_dataset(dataset, new_idxs)
    return dataset


def subsample_classes(dataset, include_classes=(0, 1, 8, 9)):

    cls_idxs = [x for x, t in enumerate(dataset.targets) if t in include_classes]

    target_xform_dict = {}
    for i, k in enumerate(include_classes):
        target_xform_dict[k] = i

    dataset = subsample_dataset(dataset, cls_idxs)

    # dataset.target_transform = lambda x: target_xform_dict[x]

    return dataset


def get_train_val_indices(train_dataset, val_split=0.2):

    train_classes = np.unique(train_dataset.targets)

    # Get train/test indices
    train_idxs = []
    val_idxs = []
    for cls in train_classes:

        cls_idxs = np.where(train_dataset.targets == cls)[0]

        v_ = np.random.choice(cls_idxs, replace=False, size=((int(val_split * len(cls_idxs))),))
        t_ = [x for x in cls_idxs if x not in v_]

        train_idxs.extend(t_)
        val_idxs.extend(v_)

    return train_idxs, val_idxs


def get_cifar_10_datasets(train_transform, test_transform, train_classes=(0, 1, 8, 9),
                       prop_train_labels=0.8, split_train_val=False, seed=0, args=None):

    np.random.seed(seed)

    # Init entire training set
    whole_training_set = CustomCIFAR10(root=cifar_10_root, transform=train_transform, train=True)

    # Get labelled training set which has subsampled classes, then subsample some indices from that
    train_dataset_labelled = subsample_classes(deepcopy(whole_training_set), include_classes=train_classes)
    subsample_indices = subsample_instances(train_dataset_labelled, prop_indices_to_subsample=prop_train_labels)
    train_dataset_labelled = subsample_dataset(train_dataset_labelled, subsample_indices)

    # Split into training and validation sets
    train_idxs, val_idxs = get_train_val_indices(train_dataset_labelled)
    train_dataset_labelled_split = subsample_dataset(deepcopy(train_dataset_labelled), train_idxs)
    val_dataset_labelled_split = subsample_dataset(deepcopy(train_dataset_labelled), val_idxs)
    val_dataset_labelled_split.transform = test_transform

    # Get unlabelled data
    unlabelled_indices = set(whole_training_set.uq_idxs) - set(train_dataset_labelled.uq_idxs)
    train_dataset_unlabelled = subsample_dataset(deepcopy(whole_training_set), np.array(list(unlabelled_indices)))

    # Sample the unlabelled data to only have overlap_cls and unlabelled_classes
    if args.n_overlap_cls != -1:
        cls2sample = list(set(args.overlap_cls).union(set(args.unlabeled_classes)))
        train_dataset_unlabelled = subsample_classes(deepcopy(train_dataset_unlabelled), include_classes=cls2sample)

    # Get test set for all classes
    test_dataset = CustomCIFAR10(root=cifar_10_root, transform=test_transform, train=False)

    # Either split train into train and val or use test set as val
    train_dataset_labelled = train_dataset_labelled_split if split_train_val else train_dataset_labelled
    val_dataset_labelled = val_dataset_labelled_split if split_train_val else None

    all_datasets = {
        'train_labelled': train_dataset_labelled,
        'train_unlabelled': train_dataset_unlabelled,
        'val': val_dataset_labelled,
        'test': test_dataset,
    }

    return all_datasets


def get_cifar_100_datasets(train_transform, test_transform, train_classes=range(80),
                       prop_train_labels=0.8, split_train_val=False, seed=0, args=None):

    np.random.seed(seed)
    coarse_labels = np.array([4, 1, 14, 8, 0, 6, 7, 7, 18, 3,
                                  3, 14, 9, 18, 7, 11, 3, 9, 7, 11,
                                  6, 11, 5, 10, 7, 6, 13, 15, 3, 15,
                                  0, 11, 1, 10, 12, 14, 16, 9, 11, 5,
                                  5, 19, 8, 8, 15, 13, 14, 17, 18, 10,
                                  16, 4, 17, 4, 2, 0, 17, 4, 18, 17,
                                  10, 3, 2, 12, 12, 16, 12, 1, 9, 19,
                                  2, 10, 0, 1, 16, 12, 9, 13, 15, 13,
                                  16, 19, 2, 4, 6, 19, 5, 5, 8, 19,
                                  18, 1, 2, 15, 6, 0, 17, 8, 14, 13])

    # Init entire training set
    whole_training_set = CustomCIFAR100(root=cifar_100_root, transform=train_transform, train=True)

    # Get labelled training set which has subsampled classes, then subsample some indices from that
    train_dataset_labelled = subsample_classes(deepcopy(whole_training_set), include_classes=train_classes)
    subsample_indices = subsample_instances(train_dataset_labelled, prop_indices_to_subsample=prop_train_labels)
    train_dataset_labelled = subsample_dataset(train_dataset_labelled, subsample_indices)

    # Split into training and validation sets
    train_idxs, val_idxs = get_train_val_indices(train_dataset_labelled)
    train_dataset_labelled_split = subsample_dataset(deepcopy(train_dataset_labelled), train_idxs)
    val_dataset_labelled_split = subsample_dataset(deepcopy(train_dataset_labelled), val_idxs)
    val_dataset_labelled_split.transform = test_transform

    # Get unlabelled data
    unlabelled_indices = set(whole_training_set.uq_idxs) - set(train_dataset_labelled.uq_idxs)
    train_dataset_unlabelled = subsample_dataset(deepcopy(whole_training_set), np.array(list(unlabelled_indices)))

    # Sample the unlabelled data to only have overlap_cls and unlabelled_classes
    if args.n_overlap_cls != -1:
        cls2sample = list(set(args.overlap_cls).union(set(args.unlabeled_classes)))
        train_dataset_unlabelled = subsample_classes(deepcopy(train_dataset_unlabelled), include_classes=cls2sample)

    # Get test set for all classes
    test_dataset = CustomCIFAR100(root=cifar_100_root, transform=test_transform, train=False)

    # Either split train into train and val or use test set as val
    train_dataset_labelled = train_dataset_labelled_split if split_train_val else train_dataset_labelled
    val_dataset_labelled = val_dataset_labelled_split if split_train_val else None

    all_datasets = {
        'train_labelled': train_dataset_labelled,
        'train_unlabelled': train_dataset_unlabelled,
        'val': val_dataset_labelled,
        'test': test_dataset,
    }

    return all_datasets

def get_cifar_100_lt_datasets(train_transform, test_transform, train_classes=range(80),
                            prop_train_labels=0.8, split_train_val=False, seed=0, args=None):
    np.random.seed(seed)
    # Init entire training set
    whole_training_set = CustomCIFAR100(root=cifar_100_root, transform=train_transform, train=True)

    # Get labelled training set which has subsampled classes, then subsample some indices from that
    train_dataset_labelled = subsample_classes(deepcopy(whole_training_set), include_classes=train_classes)
    subsample_indices = subsample_instances(train_dataset_labelled, prop_indices_to_subsample=prop_train_labels)
    train_dataset_labelled = subsample_dataset(train_dataset_labelled, subsample_indices)

    # Split into training and validation sets
    train_idxs, val_idxs = get_train_val_indices(train_dataset_labelled)
    train_dataset_labelled_split = subsample_dataset(deepcopy(train_dataset_labelled), train_idxs)
    val_dataset_labelled_split = subsample_dataset(deepcopy(train_dataset_labelled), val_idxs)
    val_dataset_labelled_split.transform = test_transform

    # Get unlabelled data
    unlabelled_indices = set(whole_training_set.uq_idxs) - set(train_dataset_labelled.uq_idxs)
    train_dataset_unlabelled = subsample_dataset(deepcopy(whole_training_set), np.array(list(unlabelled_indices)))

    # Sample the unlabelled data to only have overlap_cls and unlabelled_classes
    if args.n_overlap_cls != -1:
        cls2sample = list(set(args.overlap_cls).union(set(args.unlabeled_classes)))
        train_dataset_unlabelled = subsample_classes(deepcopy(train_dataset_unlabelled), include_classes=cls2sample)
        
    # Sample the unlabelled data to be long tail
    train_dataset_unlabelled = subsample_dataset_lt(deepcopy(train_dataset_unlabelled), args)

    # Get test set for all classes
    test_dataset = CustomCIFAR100(root=cifar_100_root, transform=test_transform, train=False)

    # Either split train into train and val or use test set as val
    train_dataset_labelled = train_dataset_labelled_split if split_train_val else train_dataset_labelled
    val_dataset_labelled = val_dataset_labelled_split if split_train_val else None

    all_datasets = {
        'train_labelled': train_dataset_labelled,
        'train_unlabelled': train_dataset_unlabelled,
        'val': val_dataset_labelled,
        'test': test_dataset,
    }

    return all_datasets


PARTITION_CONFIG_MIX = {
    # (target_list, (partition_lower, partition_upper))
    'stage 1': ((list(range(70)), (0, 0.87)),),
    'stage 2': ((list(range(70)), (0.87, 0.95)),
                (list(range(70, 80)), (0, 0.7)),),
    'stage 3': ((list(range(70)), (0.95, 0.97)),
                (list(range(70, 80)), (0.7, 0.9)),
                (list(range(80, 90)), (0, 0.9)),),
    'stage 4': ((list(range(70)), (0.97, 1.0)),
                (list(range(70, 80)), (0.9, 1.0)),
                (list(range(80, 90)), (0.9, 1.0)),
                (list(range(90, 100)), (0, 1.0)),),
}

def get_cifar_100_icd_dataset(train_transform, test_transform, train_classes=range(80),
                       prop_train_labels=0.8, split_train_val=False, stage=0, seed=0, args=None):
    # this is to return dataset used for each stage of class incremental learning
    # using the above partition config
    # train_classes, prop_train_labels, split_train_val is not used here
    np.random.seed(seed)
    config = PARTITION_CONFIG_MIX[f'stage {stage+1}']
    
    train_d, test_d = None, None
    for target_list, (lower, upper) in config:
        if train_d is None:
            train_d = PartialCIFAR100(root=cifar_100_root, split='train', transform=train_transform, download=True, target_list=target_list, partition=(lower, upper))
            test_d = PartialCIFAR100(root=cifar_100_root, split='test', transform=test_transform, download=True, target_list=target_list, partition=(lower, upper))
        else:
            train_d.concat(PartialCIFAR100(root=cifar_100_root, split='train', transform=train_transform, download=True, target_list=target_list, partition=(lower, upper)))
            test_d.concat(PartialCIFAR100(root=cifar_100_root, split='test', transform=test_transform, download=True, target_list=target_list, partition=(lower, upper)))
            
    return train_d, test_d
    


if __name__ == '__main__':

    x = get_cifar_100_datasets(None, None, split_train_val=False,
                         train_classes=range(80), prop_train_labels=0.5)

    print('Printing lens...')
    for k, v in x.items():
        if v is not None:
            print(f'{k}: {len(v)}')

    print('Printing labelled and unlabelled overlap...')
    print(set.intersection(set(x['train_labelled'].uq_idxs), set(x['train_unlabelled'].uq_idxs)))
    print('Printing total instances in train...')
    print(len(set(x['train_labelled'].uq_idxs)) + len(set(x['train_unlabelled'].uq_idxs)))

    print(f'Num Labelled Classes: {len(set(x["train_labelled"].targets))}')
    print(f'Num Unabelled Classes: {len(set(x["train_unlabelled"].targets))}')
    print(f'Len labelled set: {len(x["train_labelled"])}')
    print(f'Len unlabelled set: {len(x["train_unlabelled"])}')
data/cifar.py
DTennant/IGCD
import os
import pandas as pd
import numpy as np
from copy import deepcopy

from torchvision.datasets.folder import default_loader
from torchvision.datasets.utils import download_url
from torch.utils.data import Dataset

from data.data_utils import subsample_instances
from config import cub_root


class CustomCub2011(Dataset):
    base_folder = 'CUB_200_2011/images'
    url = 'http://www.vision.caltech.edu/visipedia-data/CUB-200-2011/CUB_200_2011.tgz'
    filename = 'CUB_200_2011.tgz'
    tgz_md5 = '97eceeb196236b17998738112f37df78'

    def __init__(self, root, train=True, transform=None, target_transform=None, loader=default_loader, download=True):

        self.root = os.path.expanduser(root)
        self.transform = transform
        self.target_transform = target_transform

        self.loader = loader
        self.train = train


        if download:
            self._download()

        if not self._check_integrity():
            raise RuntimeError('Dataset not found or corrupted.' +
                               ' You can use download=True to download it')

        self.uq_idxs = np.array(range(len(self)))

    def _load_metadata(self):
        images = pd.read_csv(os.path.join(self.root, 'CUB_200_2011', 'images.txt'), sep=' ',
                             names=['img_id', 'filepath'])
        image_class_labels = pd.read_csv(os.path.join(self.root, 'CUB_200_2011', 'image_class_labels.txt'),
                                         sep=' ', names=['img_id', 'target'])
        train_test_split = pd.read_csv(os.path.join(self.root, 'CUB_200_2011', 'train_test_split.txt'),
                                       sep=' ', names=['img_id', 'is_training_img'])

        data = images.merge(image_class_labels, on='img_id')
        self.data = data.merge(train_test_split, on='img_id')

        if self.train:
            self.data = self.data[self.data.is_training_img == 1]
        else:
            self.data = self.data[self.data.is_training_img == 0]

    def _check_integrity(self):
        try:
            self._load_metadata()
        except Exception:
            return False

        for index, row in self.data.iterrows():
            filepath = os.path.join(self.root, self.base_folder, row.filepath)
            if not os.path.isfile(filepath):
                print(filepath)
                return False
        return True

    def _download(self):
        import tarfile

        if self._check_integrity():
            print('Files already downloaded and verified')
            return

        download_url(self.url, self.root, self.filename, self.tgz_md5)

        with tarfile.open(os.path.join(self.root, self.filename), "r:gz") as tar:
            tar.extractall(path=self.root)

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        sample = self.data.iloc[idx]
        path = os.path.join(self.root, self.base_folder, sample.filepath)
        target = sample.target - 1  # Targets start at 1 by default, so shift to 0
        img = self.loader(path)

        if self.transform is not None:
            img = self.transform(img)

        if self.target_transform is not None:
            target = self.target_transform(target)

        return img, target, self.uq_idxs[idx]


def subsample_dataset(dataset, idxs):

    mask = np.zeros(len(dataset)).astype('bool')
    mask[idxs] = True

    dataset.data = dataset.data[mask]
    dataset.uq_idxs = dataset.uq_idxs[mask]

    return dataset


def subsample_classes(dataset, include_classes=range(160)):

    include_classes_cub = np.array(include_classes) + 1     # CUB classes are indexed 1 --> 200 instead of 0 --> 199
    cls_idxs = [x for x, (_, r) in enumerate(dataset.data.iterrows()) if int(r['target']) in include_classes_cub]

    # TODO: For now have no target transform
    target_xform_dict = {}
    for i, k in enumerate(include_classes):
        target_xform_dict[k] = i

    dataset = subsample_dataset(dataset, cls_idxs)

    dataset.target_transform = lambda x: target_xform_dict[x]

    return dataset


def get_train_val_indices(train_dataset, val_split=0.2):

    train_classes = np.unique(train_dataset.data['target'])

    # Get train/test indices
    train_idxs = []
    val_idxs = []
    for cls in train_classes:

        cls_idxs = np.where(train_dataset.data['target'] == cls)[0]

        v_ = np.random.choice(cls_idxs, replace=False, size=((int(val_split * len(cls_idxs))),))
        t_ = [x for x in cls_idxs if x not in v_]

        train_idxs.extend(t_)
        val_idxs.extend(v_)

    return train_idxs, val_idxs


def get_cub_datasets(train_transform, test_transform, train_classes=range(160), prop_train_labels=0.8,
                    split_train_val=False, seed=0, download=False, args=None):

    np.random.seed(seed)

    # Init entire training set
    whole_training_set = CustomCub2011(root=cub_root, transform=train_transform, train=True, download=download)

    # Get labelled training set which has subsampled classes, then subsample some indices from that
    train_dataset_labelled = subsample_classes(deepcopy(whole_training_set), include_classes=train_classes)
    subsample_indices = subsample_instances(train_dataset_labelled, prop_indices_to_subsample=prop_train_labels)
    train_dataset_labelled = subsample_dataset(train_dataset_labelled, subsample_indices)

    # Split into training and validation sets
    train_idxs, val_idxs = get_train_val_indices(train_dataset_labelled)
    train_dataset_labelled_split = subsample_dataset(deepcopy(train_dataset_labelled), train_idxs)
    val_dataset_labelled_split = subsample_dataset(deepcopy(train_dataset_labelled), val_idxs)
    val_dataset_labelled_split.transform = test_transform

    # Get unlabelled data
    unlabelled_indices = set(whole_training_set.uq_idxs) - set(train_dataset_labelled.uq_idxs)
    train_dataset_unlabelled = subsample_dataset(deepcopy(whole_training_set), np.array(list(unlabelled_indices)))
    
    # Sample the unlabelled data to only have overlap_cls and unlabelled_classes
    if args.n_overlap_cls != -1:
        cls2sample = list(set(args.overlap_cls).union(set(args.unlabeled_classes)))
        train_dataset_unlabelled = subsample_classes(deepcopy(train_dataset_unlabelled), include_classes=cls2sample)

    # Get test set for all classes
    test_dataset = CustomCub2011(root=cub_root, transform=test_transform, train=False)

    # Either split train into train and val or use test set as val
    train_dataset_labelled = train_dataset_labelled_split if split_train_val else train_dataset_labelled
    val_dataset_labelled = val_dataset_labelled_split if split_train_val else None

    all_datasets = {
        'train_labelled': train_dataset_labelled,
        'train_unlabelled': train_dataset_unlabelled,
        'val': val_dataset_labelled,
        'test': test_dataset,
    }

    return all_datasets


if __name__ == '__main__':

    x = get_cub_datasets(None, None, split_train_val=False,
                         train_classes=range(100), prop_train_labels=0.5)

    print('Printing lens...')
    for k, v in x.items():
        if v is not None:
            print(f'{k}: {len(v)}')

    print('Printing labelled and unlabelled overlap...')
    print(set.intersection(set(x['train_labelled'].uq_idxs), set(x['train_unlabelled'].uq_idxs)))
    print('Printing total instances in train...')
    print(len(set(x['train_labelled'].uq_idxs)) + len(set(x['train_unlabelled'].uq_idxs)))

    print(f'Num Labelled Classes: {len(set(x["train_labelled"].data["target"].values))}')
    print(f'Num Unabelled Classes: {len(set(x["train_unlabelled"].data["target"].values))}')
    print(f'Len labelled set: {len(x["train_labelled"])}')
    print(f'Len unlabelled set: {len(x["train_unlabelled"])}')
    

data/cub.py
DTennant/IGCD
import numpy as np
from torch.utils.data import Dataset

def subsample_instances(dataset, prop_indices_to_subsample=0.8):

    np.random.seed(0)
    subsample_indices = np.random.choice(range(len(dataset)), replace=False,
                                         size=(int(prop_indices_to_subsample * len(dataset)),))

    return subsample_indices

class MergedDataset(Dataset):

    """
    Takes two datasets (labelled_dataset, unlabelled_dataset) and merges them
    Allows you to iterate over them in parallel
    """

    def __init__(self, labelled_dataset, unlabelled_dataset):

        self.labelled_dataset = labelled_dataset
        self.unlabelled_dataset = unlabelled_dataset
        self.target_transform = None

    def __getitem__(self, item):

        if item < len(self.labelled_dataset):
            # img, label, uq_idx = self.labelled_dataset[item]
            items = self.labelled_dataset[item]
            labeled_or_not = 1

        else:

            # img, label, uq_idx = self.unlabelled_dataset[item - len(self.labelled_dataset)]
            items = self.unlabelled_dataset[item - len(self.labelled_dataset)]
            labeled_or_not = 0

        items = list(items)
        items.append(labeled_or_not)
        return tuple(items)

        # return img, label, uq_idx, np.array([labeled_or_not])

    def __len__(self):
        return len(self.unlabelled_dataset) + len(self.labelled_dataset)

data/data_utils.py
DTennant/IGCD
import os
import numpy as np
from copy import deepcopy

from torchvision.datasets.folder import default_loader
from torch.utils.data import Dataset

from data.data_utils import subsample_instances
from config import aircraft_root

def make_dataset(dir, image_ids, targets):
    assert(len(image_ids) == len(targets))
    images = []
    dir = os.path.expanduser(dir)
    for i in range(len(image_ids)):
        item = (os.path.join(dir, 'data', 'images',
                             '%s.jpg' % image_ids[i]), targets[i])
        images.append(item)
    return images


def find_classes(classes_file):

    # read classes file, separating out image IDs and class names
    image_ids = []
    targets = []
    f = open(classes_file, 'r')
    for line in f:
        split_line = line.split(' ')
        image_ids.append(split_line[0])
        targets.append(' '.join(split_line[1:]))
    f.close()

    # index class names
    classes = np.unique(targets)
    class_to_idx = {classes[i]: i for i in range(len(classes))}
    targets = [class_to_idx[c] for c in targets]

    return (image_ids, targets, classes, class_to_idx)


class FGVCAircraft(Dataset):

    """`FGVC-Aircraft <http://www.robots.ox.ac.uk/~vgg/data/fgvc-aircraft>`_ Dataset.

    Args:
        root (string): Root directory path to dataset.
        class_type (string, optional): The level of FGVC-Aircraft fine-grain classification
            to label data with (i.e., ``variant``, ``family``, or ``manufacturer``).
        transform (callable, optional): A function/transform that takes in a PIL image
            and returns a transformed version. E.g. ``transforms.RandomCrop``
        target_transform (callable, optional): A function/transform that takes in the
            target and transforms it.
        loader (callable, optional): A function to load an image given its path.
        download (bool, optional): If true, downloads the dataset from the internet and
            puts it in the root directory. If dataset is already downloaded, it is not
            downloaded again.
    """
    url = 'http://www.robots.ox.ac.uk/~vgg/data/fgvc-aircraft/archives/fgvc-aircraft-2013b.tar.gz'
    class_types = ('variant', 'family', 'manufacturer')
    splits = ('train', 'val', 'trainval', 'test')

    def __init__(self, root, class_type='variant', split='train', transform=None,
                 target_transform=None, loader=default_loader, download=False):
        if split not in self.splits:
            raise ValueError('Split "{}" not found. Valid splits are: {}'.format(
                split, ', '.join(self.splits),
            ))
        if class_type not in self.class_types:
            raise ValueError('Class type "{}" not found. Valid class types are: {}'.format(
                class_type, ', '.join(self.class_types),
            ))
        self.root = os.path.expanduser(root)
        self.class_type = class_type
        self.split = split
        self.classes_file = os.path.join(self.root, 'data',
                                         'images_%s_%s.txt' % (self.class_type, self.split))

        if download:
            self.download()

        (image_ids, targets, classes, class_to_idx) = find_classes(self.classes_file)
        samples = make_dataset(self.root, image_ids, targets)

        self.transform = transform
        self.target_transform = target_transform
        self.loader = loader

        self.samples = samples
        self.classes = classes
        self.class_to_idx = class_to_idx
        self.train = True if split == 'train' else False

        self.uq_idxs = np.array(range(len(self)))

    def __getitem__(self, index):
        """
        Args:
            index (int): Index

        Returns:
            tuple: (sample, target) where target is class_index of the target class.
        """

        path, target = self.samples[index]
        sample = self.loader(path)
        if self.transform is not None:
            sample = self.transform(sample)
        if self.target_transform is not None:
            target = self.target_transform(target)

        return sample, target, self.uq_idxs[index]

    def __len__(self):
        return len(self.samples)

    def __repr__(self):
        fmt_str = 'Dataset ' + self.__class__.__name__ + '\n'
        fmt_str += '    Number of datapoints: {}\n'.format(self.__len__())
        fmt_str += '    Root Location: {}\n'.format(self.root)
        tmp = '    Transforms (if any): '
        fmt_str += '{0}{1}\n'.format(tmp, self.transform.__repr__().replace('\n', '\n' + ' ' * len(tmp)))
        tmp = '    Target Transforms (if any): '
        fmt_str += '{0}{1}'.format(tmp, self.target_transform.__repr__().replace('\n', '\n' + ' ' * len(tmp)))
        return fmt_str

    def _check_exists(self):
        return os.path.exists(os.path.join(self.root, 'data', 'images')) and \
            os.path.exists(self.classes_file)

    def download(self):
        """Download the FGVC-Aircraft data if it doesn't exist already."""
        from six.moves import urllib
        import tarfile

        if self._check_exists():
            return

        # prepare to download data to PARENT_DIR/fgvc-aircraft-2013.tar.gz
        print('Downloading %s ... (may take a few minutes)' % self.url)
        parent_dir = os.path.abspath(os.path.join(self.root, os.pardir))
        tar_name = self.url.rpartition('/')[-1]
        tar_path = os.path.join(parent_dir, tar_name)
        data = urllib.request.urlopen(self.url)

        # download .tar.gz file
        with open(tar_path, 'wb') as f:
            f.write(data.read())

        # extract .tar.gz to PARENT_DIR/fgvc-aircraft-2013b
        data_folder = tar_path.strip('.tar.gz')
        print('Extracting %s to %s ... (may take a few minutes)' % (tar_path, data_folder))
        tar = tarfile.open(tar_path)
        tar.extractall(parent_dir)

        # if necessary, rename data folder to self.root
        if not os.path.samefile(data_folder, self.root):
            print('Renaming %s to %s ...' % (data_folder, self.root))
            os.rename(data_folder, self.root)

        # delete .tar.gz file
        print('Deleting %s ...' % tar_path)
        os.remove(tar_path)

        print('Done!')


def subsample_dataset(dataset, idxs):

    mask = np.zeros(len(dataset)).astype('bool')
    mask[idxs] = True

    dataset.samples = [(p, t) for i, (p, t) in enumerate(dataset.samples) if i in idxs]
    dataset.uq_idxs = dataset.uq_idxs[mask]

    return dataset


def subsample_classes(dataset, include_classes=range(60)):

    cls_idxs = [i for i, (p, t) in enumerate(dataset.samples) if t in include_classes]

    # TODO: Don't transform targets for now
    target_xform_dict = {}
    for i, k in enumerate(include_classes):
        target_xform_dict[k] = i

    dataset = subsample_dataset(dataset, cls_idxs)

    dataset.target_transform = lambda x: target_xform_dict[x]

    return dataset


def get_train_val_indices(train_dataset, val_split=0.2):

    all_targets = [t for i, (p, t) in enumerate(train_dataset.samples)]
    train_classes = np.unique(all_targets)

    # Get train/test indices
    train_idxs = []
    val_idxs = []
    for cls in train_classes:
        cls_idxs = np.where(all_targets == cls)[0]

        v_ = np.random.choice(cls_idxs, replace=False, size=((int(val_split * len(cls_idxs))),))
        t_ = [x for x in cls_idxs if x not in v_]

        train_idxs.extend(t_)
        val_idxs.extend(v_)

    return train_idxs, val_idxs


def get_aircraft_datasets(train_transform, test_transform, train_classes=range(50), prop_train_labels=0.8,
                    split_train_val=False, seed=0, args=None):

    np.random.seed(seed)

    # Init entire training set
    whole_training_set = FGVCAircraft(root=aircraft_root, transform=train_transform, split='trainval')

    # Get labelled training set which has subsampled classes, then subsample some indices from that
    train_dataset_labelled = subsample_classes(deepcopy(whole_training_set), include_classes=train_classes)
    subsample_indices = subsample_instances(train_dataset_labelled, prop_indices_to_subsample=prop_train_labels)
    train_dataset_labelled = subsample_dataset(train_dataset_labelled, subsample_indices)

    # Split into training and validation sets
    train_idxs, val_idxs = get_train_val_indices(train_dataset_labelled)
    train_dataset_labelled_split = subsample_dataset(deepcopy(train_dataset_labelled), train_idxs)
    val_dataset_labelled_split = subsample_dataset(deepcopy(train_dataset_labelled), val_idxs)
    val_dataset_labelled_split.transform = test_transform

    # Get unlabelled data
    unlabelled_indices = set(whole_training_set.uq_idxs) - set(train_dataset_labelled.uq_idxs)
    train_dataset_unlabelled = subsample_dataset(deepcopy(whole_training_set), np.array(list(unlabelled_indices)))

    # Sample the unlabelled data to only have overlap_cls and unlabelled_classes
    if args.n_overlap_cls != -1:
        cls2sample = list(set(args.overlap_cls).union(set(args.unlabeled_classes)))
        train_dataset_unlabelled = subsample_classes(deepcopy(train_dataset_unlabelled), include_classes=cls2sample)

    # Get test set for all classes
    test_dataset = FGVCAircraft(root=aircraft_root, transform=test_transform, split='test')

    # Either split train into train and val or use test set as val
    train_dataset_labelled = train_dataset_labelled_split if split_train_val else train_dataset_labelled
    val_dataset_labelled = val_dataset_labelled_split if split_train_val else None

    all_datasets = {
        'train_labelled': train_dataset_labelled,
        'train_unlabelled': train_dataset_unlabelled,
        'val': val_dataset_labelled,
        'test': test_dataset,
    }

    return all_datasets

if __name__ == '__main__':

    x = get_aircraft_datasets(None, None, split_train_val=False)

    print('Printing lens...')
    for k, v in x.items():
        if v is not None:
            print(f'{k}: {len(v)}')

    print('Printing labelled and unlabelled overlap...')
    print(set.intersection(set(x['train_labelled'].uq_idxs), set(x['train_unlabelled'].uq_idxs)))
    print('Printing total instances in train...')
    print(len(set(x['train_labelled'].uq_idxs)) + len(set(x['train_unlabelled'].uq_idxs)))
    print('Printing number of labelled classes...')
    print(len(set([i[1] for i in x['train_labelled'].samples])))
    print('Printing total number of classes...')
    print(len(set([i[1] for i in x['train_unlabelled'].samples])))

data/fgvc_aircraft.py
DTennant/IGCD
from data.data_utils import MergedDataset

from data.cifar import get_cifar_10_datasets, get_cifar_100_datasets, get_cifar_100_lt_datasets
from data.herbarium_19 import get_herbarium_datasets
from data.stanford_cars import get_scars_datasets
from data.imagenet import get_imagenet_100_datasets, get_imagenet_100_lt_datasets
from data.cub import get_cub_datasets
from data.fgvc_aircraft import get_aircraft_datasets
from data.nabirds import get_nabirds_datasets
from data.inat21 import get_inat21_datasets, get_inat21_incremental_datasets

from copy import deepcopy
import pickle
import os
import numpy as np

from config import osr_split_dir


get_dataset_funcs = {
    'cifar10': get_cifar_10_datasets,
    'cifar100': get_cifar_100_datasets,
    'cifar100_lt': get_cifar_100_lt_datasets,
    'cifar100_osssl': get_cifar_100_datasets,
    'imagenet_100': get_imagenet_100_datasets,
    'imagenet_100_lt': get_imagenet_100_lt_datasets,
    'herbarium_19': get_herbarium_datasets,
    'cub': get_cub_datasets,
    'aircraft': get_aircraft_datasets,
    'scars': get_scars_datasets,
    'nabirds': get_nabirds_datasets,
    'inat21': get_inat21_datasets,
    'inat21_incremental': get_inat21_incremental_datasets,
}


def get_datasets(dataset_name, train_transform, test_transform, args):

    """
    :return: train_dataset: MergedDataset which concatenates labelled and unlabelled
             test_dataset,
             unlabelled_train_examples_test,
             datasets
    """

    #
    if dataset_name not in get_dataset_funcs.keys():
        raise ValueError

    # Get datasets
    # import ipdb; ipdb.set_trace()
    get_dataset_f = get_dataset_funcs[dataset_name]
    datasets = get_dataset_f(train_transform=train_transform, test_transform=test_transform,
                            train_classes=args.train_classes,
                            prop_train_labels=args.prop_train_labels,
                            split_train_val=False,
                            args=args)
    # Set target transforms:
    # import ipdb; ipdb.set_trace();
    target_transform_dict = {}
    for i, cls in enumerate(list(args.train_classes) + list(args.unlabeled_classes)):
        target_transform_dict[cls] = i
    def _transform(x):
        return target_transform_dict[x]

    if dataset_name == 'inat21_incremental':
        def _transform(x):
            return x
            
    target_transform = _transform

    for dataset_name, dataset in datasets.items():
        if dataset is not None:
            dataset.target_transform = target_transform

    # Train split (labelled and unlabelled classes) for training
    train_dataset = MergedDataset(labelled_dataset=deepcopy(datasets['train_labelled']),
                                  unlabelled_dataset=deepcopy(datasets['train_unlabelled']))

    test_dataset = datasets['test']
    unlabelled_train_examples_test = deepcopy(datasets['train_unlabelled'])
    unlabelled_train_examples_test.transform = test_transform

    return train_dataset, test_dataset, unlabelled_train_examples_test, datasets


def get_class_splits(args):

    # For FGVC datasets, optionally return bespoke splits
    if args.dataset_name in ('scars', 'cub', 'aircraft', 'nabirds'):
        if hasattr(args, 'use_ssb_splits'):
            use_ssb_splits = args.use_ssb_splits
        else:
            use_ssb_splits = False

    # -------------
    # GET CLASS SPLITS
    # -------------
    if args.dataset_name == 'cifar10':

        args.image_size = 32
        args.train_classes = range(5)
        args.unlabeled_classes = range(5, 10)

    elif args.dataset_name == 'cifar100':

        args.image_size = 32
        args.train_classes = range(80)
        args.unlabeled_classes = range(80, 100)

    elif args.dataset_name == 'cifar100_osssl':

        args.image_size = 32
        args.train_classes = range(55)
        args.unlabeled_classes = range(55, 100)

    elif args.dataset_name == 'cifar100_lt':

        args.image_size = 32
        args.train_classes = range(80)
        args.unlabeled_classes = range(80, 100)

    elif args.dataset_name == 'herbarium_19':

        args.image_size = 224
        herb_path_splits = os.path.join(osr_split_dir, 'herbarium_19_class_splits.pkl')

        with open(herb_path_splits, 'rb') as handle:
            class_splits = pickle.load(handle)

        args.train_classes = class_splits['Old']
        args.unlabeled_classes = class_splits['New']

    elif args.dataset_name == 'imagenet_100':

        args.image_size = 224
        args.train_classes = range(50)
        args.unlabeled_classes = range(50, 100)

    elif args.dataset_name == 'imagenet_100_lt':

        args.image_size = 224
        args.train_classes = range(50)
        args.unlabeled_classes = range(50, 100)

    elif args.dataset_name == 'scars':

        args.image_size = 224

        if use_ssb_splits:

            split_path = os.path.join(osr_split_dir, 'scars_osr_splits.pkl')
            with open(split_path, 'rb') as handle:
                class_info = pickle.load(handle)

            args.train_classes = class_info['known_classes']
            open_set_classes = class_info['unknown_classes']
            args.unlabeled_classes = open_set_classes['Hard'] + open_set_classes['Medium'] + open_set_classes['Easy']
            args.hard_classes = open_set_classes['Hard']
            args.medium_classes = open_set_classes['Medium']
            args.easy_classes = open_set_classes['Easy']

        else:

            args.train_classes = range(98)
            args.unlabeled_classes = range(98, 196)

    elif args.dataset_name == 'aircraft':

        args.image_size = 224
        if use_ssb_splits:

            split_path = os.path.join(osr_split_dir, 'aircraft_osr_splits.pkl')
            with open(split_path, 'rb') as handle:
                class_info = pickle.load(handle)

            args.train_classes = class_info['known_classes']
            open_set_classes = class_info['unknown_classes']
            args.unlabeled_classes = open_set_classes['Hard'] + open_set_classes['Medium'] + open_set_classes['Easy']
            args.hard_classes = open_set_classes['Hard']
            args.medium_classes = open_set_classes['Medium']
            args.easy_classes = open_set_classes['Easy']

        else:

            args.train_classes = range(50)
            args.unlabeled_classes = range(50, 100)

    elif args.dataset_name == 'cub':

        args.image_size = 224

        if use_ssb_splits:

            split_path = os.path.join(osr_split_dir, 'cub_osr_splits.pkl')
            with open(split_path, 'rb') as handle:
                class_info = pickle.load(handle)

            args.train_classes = class_info['known_classes']
            open_set_classes = class_info['unknown_classes']
            args.unlabeled_classes = open_set_classes['Hard'] + open_set_classes['Medium'] + open_set_classes['Easy']
            args.hard_classes = open_set_classes['Hard']
            args.medium_classes = open_set_classes['Medium']
            args.easy_classes = open_set_classes['Easy']

        else:

            args.train_classes = range(100)
            args.unlabeled_classes = range(100, 200)
            
    elif args.dataset_name == 'nabirds':
        
        args.image_size = 224
        if use_ssb_splits:
            split_path = os.path.join(osr_split_dir, 'nabirds_osr_splits.pkl')
            with open(split_path, 'rb') as handle:
                class_info = pickle.load(handle)
                
            args.train_classes = class_info['Label']['id']
            args.unlabeled_classes = class_info['Hard']['id'] + class_info['Medium']['id'] + class_info['Easy']['id']

            args.hard_classes = class_info['Hard']['id']
            args.medium_classes = class_info['Medium']['id']
            args.easy_classes = class_info['Easy']['id']

        else:
            raise ValueError('Only SSB splits for NABirds')
            args.train_classes = range(200)
            args.unlabeled_classes = range(200, 555)

    elif args.dataset_name == 'inat21': 
            
        args.image_size = 224
        all_cls = np.arange(0, 10000)
        np.random.seed(42)
        args.train_classes = np.random.choice(all_cls, size=len(all_cls)//2, replace=False).tolist()
        args.unlabeled_classes = list(set(all_cls) - set(args.train_classes))

    elif args.dataset_name == 'inat21_incremental': 
            
        args.image_size = 224
        all_cls = np.arange(0, 10000)
        np.random.seed(42)
        args.train_classes = np.random.choice(all_cls, size=len(all_cls)//2, replace=False).tolist()
        args.unlabeled_classes = list(set(all_cls) - set(args.train_classes))
        
    elif args.dataset_name == 'inat18':
        args.image_size = 224
        all_cls = np.arange(0, 8142)
        np.random.seed(42)
        args.train_classes = np.random.choice(all_cls, size=len(all_cls)//2, replace=False).tolist()
        args.unlabeled_classes = list(set(all_cls) - set(args.train_classes))

    elif args.dataset_name == 'inat19':
        #TODO: SSB splits
        args.image_size = 224
        all_cls = np.arange(0, 1010)
        np.random.seed(42)
        args.train_classes = np.random.choice(all_cls, size=len(all_cls)//2, replace=False).tolist()
        args.unlabeled_classes = list(set(all_cls) - set(args.train_classes))


    else:

        raise NotImplementedError

    if args.n_overlap_cls != -1:
        # move args.n_overlap_cls classes from train to unlabeled
        np.random.seed(42)
        args.overlap_cls = np.random.choice(args.train_classes, size=args.n_overlap_cls, replace=False).tolist()
    else:
        args.overlap_cls = []

    return args

data/get_datasets.py
DTennant/IGCD
import os

import torchvision
import numpy as np
from copy import deepcopy

from data.data_utils import subsample_instances
from config import herbarium_dataroot

class HerbariumDataset19(torchvision.datasets.ImageFolder):

    def __init__(self, *args, **kwargs):

        # Process metadata json for training images into a DataFrame
        super().__init__(*args, **kwargs)

        self.uq_idxs = np.array(range(len(self)))

    def __getitem__(self, idx):

        img, label = super().__getitem__(idx)
        uq_idx = self.uq_idxs[idx]

        return img, label, uq_idx


def subsample_dataset(dataset, idxs):

    mask = np.zeros(len(dataset)).astype('bool')
    mask[idxs] = True

    dataset.samples = np.array(dataset.samples)[mask].tolist()
    dataset.targets = np.array(dataset.targets)[mask].tolist()

    dataset.uq_idxs = dataset.uq_idxs[mask]

    dataset.samples = [[x[0], int(x[1])] for x in dataset.samples]
    dataset.targets = [int(x) for x in dataset.targets]

    return dataset


def subsample_classes(dataset, include_classes=range(250)):

    cls_idxs = [x for x, l in enumerate(dataset.targets) if l in include_classes]

    target_xform_dict = {}
    for i, k in enumerate(include_classes):
        target_xform_dict[k] = i

    dataset = subsample_dataset(dataset, cls_idxs)

    dataset.target_transform = lambda x: target_xform_dict[x]

    return dataset


def get_train_val_indices(train_dataset, val_instances_per_class=5):

    train_classes = list(set(train_dataset.targets))

    # Get train/test indices
    train_idxs = []
    val_idxs = []
    for cls in train_classes:

        cls_idxs = np.where(np.array(train_dataset.targets) == cls)[0]

        # Have a balanced test set
        v_ = np.random.choice(cls_idxs, replace=False, size=(val_instances_per_class,))
        t_ = [x for x in cls_idxs if x not in v_]

        train_idxs.extend(t_)
        val_idxs.extend(v_)

    return train_idxs, val_idxs


def get_herbarium_datasets(train_transform, test_transform, train_classes=range(500), prop_train_labels=0.8,
                            seed=0, split_train_val=False, args=None):

    np.random.seed(seed)

    # Init entire training set
    train_dataset = HerbariumDataset19(transform=train_transform,
                                            root=os.path.join(herbarium_dataroot, 'small-train'))

    # Get labelled training set which has subsampled classes, then subsample some indices from that
    # TODO: Subsampling unlabelled set in uniform random fashion from training data, will contain many instances of dominant class
    train_dataset_labelled = subsample_classes(deepcopy(train_dataset), include_classes=train_classes)
    subsample_indices = subsample_instances(train_dataset_labelled, prop_indices_to_subsample=prop_train_labels)
    train_dataset_labelled = subsample_dataset(train_dataset_labelled, subsample_indices)

    # Split into training and validation sets
    if split_train_val:

        train_idxs, val_idxs = get_train_val_indices(train_dataset_labelled,
                                                     val_instances_per_class=5)
        train_dataset_labelled_split = subsample_dataset(deepcopy(train_dataset_labelled), train_idxs)
        val_dataset_labelled_split = subsample_dataset(deepcopy(train_dataset_labelled), val_idxs)
        val_dataset_labelled_split.transform = test_transform

    else:

        train_dataset_labelled_split, val_dataset_labelled_split = None, None

    # Get unlabelled data
    unlabelled_indices = set(train_dataset.uq_idxs) - set(train_dataset_labelled.uq_idxs)
    train_dataset_unlabelled = subsample_dataset(deepcopy(train_dataset), np.array(list(unlabelled_indices)))

    # Sample the unlabelled data to only have overlap_cls and unlabelled_classes
    if args.n_overlap_cls != -1:
        cls2sample = list(set(args.overlap_cls).union(set(args.unlabeled_classes)))
        train_dataset_unlabelled = subsample_classes(deepcopy(train_dataset_unlabelled), include_classes=cls2sample)

    # Get test dataset
    test_dataset = HerbariumDataset19(transform=test_transform,
                                            root=os.path.join(herbarium_dataroot, 'small-validation'))

    # Transform dict
    unlabelled_classes = list(set(train_dataset.targets) - set(train_classes))
    target_xform_dict = {}
    for i, k in enumerate(list(train_classes) + unlabelled_classes):
        target_xform_dict[k] = i

    test_dataset.target_transform = lambda x: target_xform_dict[x]
    train_dataset_unlabelled.target_transform = lambda x: target_xform_dict[x]

    # Either split train into train and val or use test set as val
    train_dataset_labelled = train_dataset_labelled_split if split_train_val else train_dataset_labelled
    val_dataset_labelled = val_dataset_labelled_split if split_train_val else None

    all_datasets = {
        'train_labelled': train_dataset_labelled,
        'train_unlabelled': train_dataset_unlabelled,
        'val': val_dataset_labelled,
        'test': test_dataset,
    }

    return all_datasets

if __name__ == '__main__':

    np.random.seed(0)
    train_classes = np.random.choice(range(683,), size=(int(683 / 2)), replace=False)

    x = get_herbarium_datasets(None, None, train_classes=train_classes,
                               prop_train_labels=0.5)

    assert set(x['train_unlabelled'].targets) == set(range(683))

    print('Printing lens...')
    for k, v in x.items():
        if v is not None:
            print(f'{k}: {len(v)}')

    print('Printing labelled and unlabelled overlap...')
    print(set.intersection(set(x['train_labelled'].uq_idxs), set(x['train_unlabelled'].uq_idxs)))
    print('Printing total instances in train...')
    print(len(set(x['train_labelled'].uq_idxs)) + len(set(x['train_unlabelled'].uq_idxs)))
    print('Printing number of labelled classes...')
    print(len(set(x['train_labelled'].targets)))
    print('Printing total number of classes...')
    print(len(set(x['train_unlabelled'].targets)))

    print(f'Num Labelled Classes: {len(set(x["train_labelled"].targets))}')
    print(f'Num Unabelled Classes: {len(set(x["train_unlabelled"].targets))}')
    print(f'Len labelled set: {len(x["train_labelled"])}')
    print(f'Len unlabelled set: {len(x["train_unlabelled"])}')
data/herbarium_19.py
DTennant/IGCD
import torchvision
import numpy as np

import os

from copy import deepcopy
from data.data_utils import subsample_instances
from config import imagenet_root


class ImageNetBase(torchvision.datasets.ImageFolder):

    def __init__(self, root, transform):

        super(ImageNetBase, self).__init__(root, transform)

        self.uq_idxs = np.array(range(len(self)))

    def __getitem__(self, item):

        img, label = super().__getitem__(item)
        uq_idx = self.uq_idxs[item]

        return img, label, uq_idx


def subsample_dataset(dataset, idxs):

    imgs_ = []
    for i in idxs:
        imgs_.append(dataset.imgs[i])
    dataset.imgs = imgs_

    samples_ = []
    for i in idxs:
        samples_.append(dataset.samples[i])
    dataset.samples = samples_

    # dataset.imgs = [x for i, x in enumerate(dataset.imgs) if i in idxs]
    # dataset.samples = [x for i, x in enumerate(dataset.samples) if i in idxs]

    dataset.targets = np.array(dataset.targets)[idxs].tolist()
    dataset.uq_idxs = dataset.uq_idxs[idxs]

    return dataset

def subsample_dataset_lt(dataset, args=None):
    cls_num = len(np.unique(np.array(dataset.targets)))
    def get_img_num_per_cls(dataset, imb_type='exp', imb_factor=0.02):
        img_max = len(dataset) / cls_num
        img_num_per_cls = []
        if imb_type == 'exp':
            for cls_idx in range(cls_num):
                num = img_max * (imb_factor**(cls_idx / (cls_num - 1.0)))
                img_num_per_cls.append(int(num))
        elif imb_type == 'step':
            for cls_idx in range(cls_num // 2):
                img_num_per_cls.append(int(img_max))
            for cls_idx in range(cls_num // 2):
                img_num_per_cls.append(int(img_max * imb_factor))
        else:
            img_num_per_cls.extend([int(img_max)] * cls_num)
        return img_num_per_cls
    
    def gen_imbalanced_data(dataset, img_num_per_cls):
        new_idxs = []
        targets_np = np.array(dataset.targets, dtype=np.int64)
        classes = np.unique(targets_np)
        dataset.num_per_cls_dict = dict()
        for the_class, the_img_num in zip(classes, img_num_per_cls):
            dataset.num_per_cls_dict[the_class] = the_img_num
            idx = np.where(targets_np == the_class)[0]
            np.random.shuffle(idx)
            selec_idx = idx[:the_img_num]
            new_idxs.append(selec_idx)
        new_idxs = np.concatenate(new_idxs, axis=0)
        return new_idxs

    img_num_per_cls = get_img_num_per_cls(dataset, imb_type='exp', imb_factor=args.imbalance_factor)
    new_idxs = gen_imbalanced_data(dataset, img_num_per_cls)
    dataset = subsample_dataset(dataset, new_idxs)
    return dataset


def subsample_classes(dataset, include_classes=list(range(1000))):

    cls_idxs = [x for x, t in enumerate(dataset.targets) if t in include_classes]

    target_xform_dict = {}
    for i, k in enumerate(include_classes):
        target_xform_dict[k] = i

    dataset = subsample_dataset(dataset, cls_idxs)
    dataset.target_transform = lambda x: target_xform_dict[x]

    return dataset


def get_train_val_indices(train_dataset, val_split=0.2):

    train_classes = list(set(train_dataset.targets))

    # Get train/test indices
    train_idxs = []
    val_idxs = []
    for cls in train_classes:

        cls_idxs = np.where(np.array(train_dataset.targets) == cls)[0]

        v_ = np.random.choice(cls_idxs, replace=False, size=((int(val_split * len(cls_idxs))),))
        t_ = [x for x in cls_idxs if x not in v_]

        train_idxs.extend(t_)
        val_idxs.extend(v_)

    return train_idxs, val_idxs


def get_imagenet_100_datasets(train_transform, test_transform, train_classes=range(80),
                           prop_train_labels=0.8, split_train_val=False, seed=0, args=None):

    np.random.seed(seed)

    # Subsample imagenet dataset initially to include 100 classes
    subsampled_100_classes = np.random.choice(range(1000), size=(100,), replace=False)
    subsampled_100_classes = np.sort(subsampled_100_classes)
    print(f'Constructing ImageNet-100 dataset from the following classes: {subsampled_100_classes.tolist()}')
    cls_map = {i: j for i, j in zip(subsampled_100_classes, range(100))}

    # Init entire training set
    imagenet_training_set = ImageNetBase(root=os.path.join(imagenet_root, 'train'), transform=train_transform)
    whole_training_set = subsample_classes(imagenet_training_set, include_classes=subsampled_100_classes)

    # Reset dataset
    whole_training_set.samples = [(s[0], cls_map[s[1]]) for s in whole_training_set.samples]
    whole_training_set.targets = [s[1] for s in whole_training_set.samples]
    whole_training_set.uq_idxs = np.array(range(len(whole_training_set)))
    whole_training_set.target_transform = None

    # Get labelled training set which has subsampled classes, then subsample some indices from that
    train_dataset_labelled = subsample_classes(deepcopy(whole_training_set), include_classes=train_classes)
    subsample_indices = subsample_instances(train_dataset_labelled, prop_indices_to_subsample=prop_train_labels)
    train_dataset_labelled = subsample_dataset(train_dataset_labelled, subsample_indices)

    # Split into training and validation sets
    train_idxs, val_idxs = get_train_val_indices(train_dataset_labelled)
    train_dataset_labelled_split = subsample_dataset(deepcopy(train_dataset_labelled), train_idxs)
    val_dataset_labelled_split = subsample_dataset(deepcopy(train_dataset_labelled), val_idxs)
    val_dataset_labelled_split.transform = test_transform

    # Get unlabelled data
    unlabelled_indices = set(whole_training_set.uq_idxs) - set(train_dataset_labelled.uq_idxs)
    train_dataset_unlabelled = subsample_dataset(deepcopy(whole_training_set), np.array(list(unlabelled_indices)))

    # Sample the unlabelled data to only have overlap_cls and unlabelled_classes
    if args.n_overlap_cls != -1:
        cls2sample = list(set(args.overlap_cls).union(set(args.unlabeled_classes)))
        train_dataset_unlabelled = subsample_classes(deepcopy(train_dataset_unlabelled), include_classes=cls2sample)

    # Get test set for all classes
    test_dataset = ImageNetBase(root=os.path.join(imagenet_root, 'val'), transform=test_transform)
    test_dataset = subsample_classes(test_dataset, include_classes=subsampled_100_classes)

    # Reset test set
    test_dataset.samples = [(s[0], cls_map[s[1]]) for s in test_dataset.samples]
    test_dataset.targets = [s[1] for s in test_dataset.samples]
    test_dataset.uq_idxs = np.array(range(len(test_dataset)))
    test_dataset.target_transform = None

    # Either split train into train and val or use test set as val
    train_dataset_labelled = train_dataset_labelled_split if split_train_val else train_dataset_labelled
    val_dataset_labelled = val_dataset_labelled_split if split_train_val else None

    all_datasets = {
        'train_labelled': train_dataset_labelled,
        'train_unlabelled': train_dataset_unlabelled,
        'val': val_dataset_labelled,
        'test': test_dataset,
    }

    return all_datasets

def get_imagenet_100_lt_datasets(train_transform, test_transform, train_classes=range(80),
                           prop_train_labels=0.8, split_train_val=False, seed=0, args=None):

    np.random.seed(seed)

    # Subsample imagenet dataset initially to include 100 classes
    subsampled_100_classes = np.random.choice(range(1000), size=(100,), replace=False)
    subsampled_100_classes = np.sort(subsampled_100_classes)
    print(f'Constructing ImageNet-100 dataset from the following classes: {subsampled_100_classes.tolist()}')
    cls_map = {i: j for i, j in zip(subsampled_100_classes, range(100))}

    # Init entire training set
    imagenet_training_set = ImageNetBase(root=os.path.join(imagenet_root, 'train'), transform=train_transform)
    whole_training_set = subsample_classes(imagenet_training_set, include_classes=subsampled_100_classes)

    # Reset dataset
    whole_training_set.samples = [(s[0], cls_map[s[1]]) for s in whole_training_set.samples]
    whole_training_set.targets = [s[1] for s in whole_training_set.samples]
    whole_training_set.uq_idxs = np.array(range(len(whole_training_set)))
    whole_training_set.target_transform = None

    # Get labelled training set which has subsampled classes, then subsample some indices from that
    train_dataset_labelled = subsample_classes(deepcopy(whole_training_set), include_classes=train_classes)
    subsample_indices = subsample_instances(train_dataset_labelled, prop_indices_to_subsample=prop_train_labels)
    train_dataset_labelled = subsample_dataset(train_dataset_labelled, subsample_indices)

    # Split into training and validation sets
    train_idxs, val_idxs = get_train_val_indices(train_dataset_labelled)
    train_dataset_labelled_split = subsample_dataset(deepcopy(train_dataset_labelled), train_idxs)
    val_dataset_labelled_split = subsample_dataset(deepcopy(train_dataset_labelled), val_idxs)
    val_dataset_labelled_split.transform = test_transform

    # Get unlabelled data
    unlabelled_indices = set(whole_training_set.uq_idxs) - set(train_dataset_labelled.uq_idxs)
    train_dataset_unlabelled = subsample_dataset(deepcopy(whole_training_set), np.array(list(unlabelled_indices)))
    
    train_dataset_unlabelled = subsample_dataset_lt(deepcopy(train_dataset_unlabelled), args)

    # Sample the unlabelled data to only have overlap_cls and unlabelled_classes
    if args.n_overlap_cls != -1:
        cls2sample = list(set(args.overlap_cls).union(set(args.unlabeled_classes)))
        train_dataset_unlabelled = subsample_classes(deepcopy(train_dataset_unlabelled), include_classes=cls2sample)

    # Get test set for all classes
    test_dataset = ImageNetBase(root=os.path.join(imagenet_root, 'val'), transform=test_transform)
    test_dataset = subsample_classes(test_dataset, include_classes=subsampled_100_classes)

    # Reset test set
    test_dataset.samples = [(s[0], cls_map[s[1]]) for s in test_dataset.samples]
    test_dataset.targets = [s[1] for s in test_dataset.samples]
    test_dataset.uq_idxs = np.array(range(len(test_dataset)))
    test_dataset.target_transform = None

    # Either split train into train and val or use test set as val
    train_dataset_labelled = train_dataset_labelled_split if split_train_val else train_dataset_labelled
    val_dataset_labelled = val_dataset_labelled_split if split_train_val else None

    all_datasets = {
        'train_labelled': train_dataset_labelled,
        'train_unlabelled': train_dataset_unlabelled,
        'val': val_dataset_labelled,
        'test': test_dataset,
    }
    return all_datasets

if __name__ == '__main__':

    x = get_imagenet_100_datasets(None, None, split_train_val=False,
                               train_classes=range(50), prop_train_labels=0.5)

    print('Printing lens...')
    for k, v in x.items():
        if v is not None:
            print(f'{k}: {len(v)}')

    print('Printing labelled and unlabelled overlap...')
    print(set.intersection(set(x['train_labelled'].uq_idxs), set(x['train_unlabelled'].uq_idxs)))
    print('Printing total instances in train...')
    print(len(set(x['train_labelled'].uq_idxs)) + len(set(x['train_unlabelled'].uq_idxs)))

    print(f'Num Labelled Classes: {len(set(x["train_labelled"].targets))}')
    print(f'Num Unabelled Classes: {len(set(x["train_unlabelled"].targets))}')
    print(f'Len labelled set: {len(x["train_labelled"])}')
    print(f'Len unlabelled set: {len(x["train_unlabelled"])}')
data/imagenet.py
DTennant/IGCD
import os
import pandas as pd
import numpy as np
from copy import deepcopy

from torchvision.datasets.folder import default_loader
from torchvision.datasets.utils import download_url
from torch.utils.data import Dataset
from sklearn.model_selection import train_test_split

from data.data_utils import subsample_instances
from config import inat21_root


class iNat(Dataset):

    def __init__(self, root, train=True, transform=None, target_transform=None, loader=default_loader, download=False):

        self.root = os.path.expanduser(root)
        self.transform = transform
        self.target_transform = target_transform

        self.loader = loader
        self.train = train

        if not self._check_integrity():
            raise RuntimeError('Dataset not found or corrupted.')

        self.uq_idxs = np.array(range(len(self)))

    def _load_metadata(self):

        self.data = pd.read_csv(os.path.join(self.root, 'train_val_mini.csv'), 
                              names=['img_id', 'filepath', 'target', 'cls_name', 'is_training_img'])

        if self.train:
            self.data = self.data[self.data.is_training_img == 1]
        else:
            self.data = self.data[self.data.is_training_img == 0]

    def _check_integrity(self):
        try:
            self._load_metadata()
        except Exception:
            return False

        # for index, row in self.data.iterrows():
        #     filepath = os.path.join(self.root, row.filepath)
        #     if not os.path.isfile(filepath):
        #         print(filepath)
        #         return False

        return True


    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        sample = self.data.iloc[idx]
        path = os.path.join(self.root, sample.filepath)
        target = int(sample.target)
        img = self.loader(path)

        if self.transform is not None:
            img = self.transform(img)

        if self.target_transform is not None:
            try:
                target = self.target_transform(target)
            except KeyError:
                # import ipdb; ipdb.set_trace()
                print('target', target)
                print('idx', idx)
                raise ValueError

        return img, target, self.uq_idxs[idx]

class iNat_incremental(Dataset):

    def __init__(self, root, use_small=True, split_crit='location', stage=1, 
                return_path=False, train=True, transform=None, target_transform=None, 
                loader=default_loader, download=False):
        # stage indicates the stage of incremental learning
        # split_crit indicates the split criterion: can be location, year, or loc_year
        # use_small indicates whether to use the small dataset

        self.root = os.path.expanduser(root)
        self.stage = stage
        self.use_small = use_small
        self.split_crit = split_crit
        self.return_path = return_path
        self.transform = transform
        self.target_transform = target_transform

        self.loader = loader
        self.train = train

        if not self._check_integrity():
            raise RuntimeError('Dataset not found or corrupted.')

        self.uq_idxs = np.array(range(len(self)))

    def _load_metadata(self):

        # self.data = pd.read_csv(os.path.join(self.root, 'train_val_mini.csv'), )
                              # names=['img_id', 'filepath', 'target', 'cls_name', 'is_training_img'])

        if self.train:
            self.data = pd.read_csv(os.path.join(self.root, 'incremental_splits', self.split_crit + '_small' if self.use_small else self.split_crit, f'stage_{self.stage}.csv'))
            # self.data = self.data[['image_id', 'category_id', 'file_name', 'date', 'continent', 'common_name', 'genus']]
        else:
            self.data = pd.read_csv(os.path.join(self.root, 'incremental_splits', self.split_crit + '_small', f'val_stage_{self.stage}.csv'))
            # self.data = self.data[['image_id', 'category_id', 'file_name', 'date', 'continent', 'common_name', 'genus']]
            
        self.data = self.data[['file_name', 'category_id']]

    def _check_integrity(self):
        try:
            self._load_metadata()
        except Exception:
            return False
        return True

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        sample = self.data.iloc[idx]
        path = os.path.join(self.root, sample.file_name)
        target = int(sample.category_id)
        img = self.loader(path)

        if self.transform is not None:
            img = self.transform(img)

        if self.target_transform is not None:
            try:
                target = self.target_transform(target)
            except KeyError:
                # import ipdb; ipdb.set_trace()
                print('target', target)
                print('idx', idx)
                raise ValueError

        if self.return_path:
            return img, target, self.uq_idxs[idx], sample.file_name
        return img, target, self.uq_idxs[idx]


def subsample_dataset(dataset, idxs):

    mask = np.zeros(len(dataset)).astype('bool')
    mask[idxs] = True

    dataset.data = dataset.data[mask]
    dataset.uq_idxs = dataset.uq_idxs[mask]

    return dataset


def subsample_classes(dataset, include_classes=range(160)):

    include_classes_cub = np.array(include_classes)
    cls_idxs = [x for x, (_, r) in enumerate(dataset.data.iterrows()) if int(r['target']) in include_classes_cub]

    target_xform_dict = {}
    for i, k in enumerate(include_classes):
        target_xform_dict[k] = i

    dataset = subsample_dataset(dataset, cls_idxs)

    dataset.target_transform = lambda x: target_xform_dict[x]

    return dataset


def get_train_val_indices(train_dataset, val_split=0.2):

    train_classes = np.unique(train_dataset.data['target'])

    # Get train/test indices
    train_idxs = []
    val_idxs = []
    for cls in train_classes:

        cls_idxs = np.where(train_dataset.data['target'] == cls)[0]

        v_ = np.random.choice(cls_idxs, replace=False, size=((int(val_split * len(cls_idxs))),))
        t_ = [x for x in cls_idxs if x not in v_]

        train_idxs.extend(t_)
        val_idxs.extend(v_)

    return train_idxs, val_idxs


def get_inat21_datasets(train_transform, test_transform, train_classes=range(160), prop_train_labels=0.8,
                    split_train_val=False, seed=0, args=None):
    # import ipdb; ipdb.set_trace()
    np.random.seed(seed)

    # Init entire training set
    whole_training_set = iNat(root=inat21_root, transform=train_transform, train=True, )

    # Get labelled training set which has subsampled classes, then subsample some indices from that
    train_dataset_labelled = subsample_classes(deepcopy(whole_training_set), include_classes=train_classes)
    subsample_indices = subsample_instances(train_dataset_labelled, prop_indices_to_subsample=prop_train_labels)
    train_dataset_labelled = subsample_dataset(train_dataset_labelled, subsample_indices)

    # Split into training and validation sets
    train_idxs, val_idxs = get_train_val_indices(train_dataset_labelled)
    train_dataset_labelled_split = subsample_dataset(deepcopy(train_dataset_labelled), train_idxs)
    val_dataset_labelled_split = subsample_dataset(deepcopy(train_dataset_labelled), val_idxs)
    val_dataset_labelled_split.transform = test_transform

    # Get unlabelled data
    unlabelled_indices = set(whole_training_set.uq_idxs) - set(train_dataset_labelled.uq_idxs)
    train_dataset_unlabelled = subsample_dataset(deepcopy(whole_training_set), np.array(list(unlabelled_indices)))

    # Sample the unlabelled data to only have overlap_cls and unlabelled_classes
    if args.n_overlap_cls != -1:
        cls2sample = list(set(args.overlap_cls).union(set(args.unlabeled_classes)))
        train_dataset_unlabelled = subsample_classes(deepcopy(train_dataset_unlabelled), include_classes=cls2sample)

    # Get test set for all classes
    test_dataset = iNat(root=inat21_root, transform=test_transform, train=False)

    # Either split train into train and val or use test set as val
    train_dataset_labelled = train_dataset_labelled_split if split_train_val else train_dataset_labelled
    val_dataset_labelled = val_dataset_labelled_split if split_train_val else None

    all_datasets = {
        'train_labelled': train_dataset_labelled,
        'train_unlabelled': train_dataset_unlabelled,
        'val': val_dataset_labelled,
        'test': test_dataset,
    }

    return all_datasets

def get_inat21_incremental_datasets(train_transform, test_transform, train_classes=range(160), prop_train_labels=0.8,
                    split_train_val=False, seed=0, args=None):
    # import ipdb; ipdb.set_trace()
    # train_classes, prop_train_labels, split_train_val, seed will be ignored
    np.random.seed(seed)

    # Init entire training set
    whole_training_set = iNat_incremental(root=inat21_root, use_small=args.use_small_set, split_crit=args.split_crit, stage=args.stage, return_path=args.return_path, transform=train_transform, train=True, )

    next_stage = args.stage + 1 if args.stage < args.max_stage else args.stage
    train_unlabelled = iNat_incremental(root=inat21_root, use_small=args.use_small_set, split_crit=args.split_crit, stage=next_stage, 
                                    transform=train_transform, train=True, return_path=args.return_path)
    test_dataset = iNat_incremental(root=inat21_root, use_small=args.use_small_set, split_crit=args.split_crit, stage=next_stage, return_path=args.return_path, transform=test_transform, train=False)

    args.classes_in_labelled = set(whole_training_set.data.category_id)
    args.classes_in_unlabelled = set(train_unlabelled.data.category_id) - set(args.classes_in_labelled)
    args.classes_in_labelled = args.classes_in_labelled - args.classes_in_unlabelled
    
    args.train_classes = args.classes_in_labelled
    args.unlabeled_classes = args.classes_in_unlabelled

    all_datasets = {
        'train_labelled': whole_training_set,
        'train_unlabelled': train_unlabelled,
        'val': None,
        'test': test_dataset,
    }

    return all_datasets


if __name__ == '__main__':

    x = get_cub_datasets(None, None, split_train_val=False,
                         train_classes=range(100), prop_train_labels=0.5)

    print('Printing lens...')
    for k, v in x.items():
        if v is not None:
            print(f'{k}: {len(v)}')

    print('Printing labelled and unlabelled overlap...')
    print(set.intersection(set(x['train_labelled'].uq_idxs), set(x['train_unlabelled'].uq_idxs)))
    print('Printing total instances in train...')
    print(len(set(x['train_labelled'].uq_idxs)) + len(set(x['train_unlabelled'].uq_idxs)))

    print(f'Num Labelled Classes: {len(set(x["train_labelled"].data["target"].values))}')
    print(f'Num Unabelled Classes: {len(set(x["train_unlabelled"].data["target"].values))}')
    print(f'Len labelled set: {len(x["train_labelled"])}')
    print(f'Len unlabelled set: {len(x["train_unlabelled"])}')
    

data/inat21.py
DTennant/IGCD
import os
import pandas as pd
import numpy as np
from copy import deepcopy

from torchvision.datasets.folder import default_loader
from torchvision.datasets.utils import download_url
from torch.utils.data import Dataset

from data.data_utils import subsample_instances
from config import nabirds_root


class NABirds(Dataset):
    base_folder = 'nabirds/images'

    def __init__(self, root, train=True, transform=None, target_transform=None, loader=default_loader, download=False):

        self.root = os.path.expanduser(root)
        self.transform = transform
        self.target_transform = target_transform

        self.loader = loader
        self.train = train

        if not self._check_integrity():
            raise RuntimeError('Dataset not found or corrupted.' +
                               ' You can download it from https://dl.allaboutbirds.org/nabirds, and also move ssb_processed.csv to the nabirds folder')

        self.uq_idxs = np.array(range(len(self)))

    def _load_metadata(self):

        self.data = pd.read_csv(os.path.join(self.root, 'nabirds', 'ssb_processed.csv'), names=['img_id', 'filepath', 'is_training_img', 'target', 'cls_name'])

        if self.train:
            self.data = self.data[self.data.is_training_img == 1]
        else:
            self.data = self.data[self.data.is_training_img == 0]

    def _check_integrity(self):
        try:
            self._load_metadata()
        except Exception:
            return False

        for index, row in self.data.iterrows():
            filepath = os.path.join(self.root, self.base_folder, row.filepath)
            if not os.path.isfile(filepath):
                print(filepath)
                return False
        return True


    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        sample = self.data.iloc[idx]
        path = os.path.join(self.root, self.base_folder, sample.filepath)
        target = int(sample.target)
        img = self.loader(path)

        if self.transform is not None:
            img = self.transform(img)

        if self.target_transform is not None:
            try:
                target = self.target_transform(target)
            except KeyError:
                # import ipdb; ipdb.set_trace()
                print('target', target)
                print('idx', idx)
                raise ValueError

        return img, target, self.uq_idxs[idx]


def subsample_dataset(dataset, idxs):

    mask = np.zeros(len(dataset)).astype('bool')
    mask[idxs] = True

    dataset.data = dataset.data[mask]
    dataset.uq_idxs = dataset.uq_idxs[mask]

    return dataset


def subsample_classes(dataset, include_classes=range(160)):

    include_classes_cub = np.array(include_classes)
    cls_idxs = [x for x, (_, r) in enumerate(dataset.data.iterrows()) if int(r['target']) in include_classes_cub]

    target_xform_dict = {}
    for i, k in enumerate(include_classes):
        target_xform_dict[k] = i

    dataset = subsample_dataset(dataset, cls_idxs)

    dataset.target_transform = lambda x: target_xform_dict[x]

    return dataset


def get_train_val_indices(train_dataset, val_split=0.2):

    train_classes = np.unique(train_dataset.data['target'])

    # Get train/test indices
    train_idxs = []
    val_idxs = []
    for cls in train_classes:

        cls_idxs = np.where(train_dataset.data['target'] == cls)[0]

        v_ = np.random.choice(cls_idxs, replace=False, size=((int(val_split * len(cls_idxs))),))
        t_ = [x for x in cls_idxs if x not in v_]

        train_idxs.extend(t_)
        val_idxs.extend(v_)

    return train_idxs, val_idxs


def get_nabirds_datasets(train_transform, test_transform, train_classes=range(160), prop_train_labels=0.8,
                    split_train_val=False, seed=0, args=None):

    np.random.seed(seed)

    # Init entire training set
    whole_training_set = NABirds(root=nabirds_root, transform=train_transform, train=True, )

    # Get labelled training set which has subsampled classes, then subsample some indices from that
    train_dataset_labelled = subsample_classes(deepcopy(whole_training_set), include_classes=train_classes)
    subsample_indices = subsample_instances(train_dataset_labelled, prop_indices_to_subsample=prop_train_labels)
    train_dataset_labelled = subsample_dataset(train_dataset_labelled, subsample_indices)

    # Split into training and validation sets
    train_idxs, val_idxs = get_train_val_indices(train_dataset_labelled)
    train_dataset_labelled_split = subsample_dataset(deepcopy(train_dataset_labelled), train_idxs)
    val_dataset_labelled_split = subsample_dataset(deepcopy(train_dataset_labelled), val_idxs)
    val_dataset_labelled_split.transform = test_transform

    # Get unlabelled data
    unlabelled_indices = set(whole_training_set.uq_idxs) - set(train_dataset_labelled.uq_idxs)
    train_dataset_unlabelled = subsample_dataset(deepcopy(whole_training_set), np.array(list(unlabelled_indices)))

    # Sample the unlabelled data to only have overlap_cls and unlabelled_classes
    if args.n_overlap_cls != -1:
        cls2sample = list(set(args.overlap_cls).union(set(args.unlabeled_classes)))
        train_dataset_unlabelled = subsample_classes(deepcopy(train_dataset_unlabelled), include_classes=cls2sample)

    # Get test set for all classes
    test_dataset = NABirds(root=nabirds_root, transform=test_transform, train=False)

    # Either split train into train and val or use test set as val
    train_dataset_labelled = train_dataset_labelled_split if split_train_val else train_dataset_labelled
    val_dataset_labelled = val_dataset_labelled_split if split_train_val else None

    all_datasets = {
        'train_labelled': train_dataset_labelled,
        'train_unlabelled': train_dataset_unlabelled,
        'val': val_dataset_labelled,
        'test': test_dataset,
    }

    return all_datasets

if __name__ == '__main__':

    x = get_cub_datasets(None, None, split_train_val=False,
                         train_classes=range(100), prop_train_labels=0.5)

    print('Printing lens...')
    for k, v in x.items():
        if v is not None:
            print(f'{k}: {len(v)}')

    print('Printing labelled and unlabelled overlap...')
    print(set.intersection(set(x['train_labelled'].uq_idxs), set(x['train_unlabelled'].uq_idxs)))
    print('Printing total instances in train...')
    print(len(set(x['train_labelled'].uq_idxs)) + len(set(x['train_unlabelled'].uq_idxs)))

    print(f'Num Labelled Classes: {len(set(x["train_labelled"].data["target"].values))}')
    print(f'Num Unabelled Classes: {len(set(x["train_unlabelled"].data["target"].values))}')
    print(f'Len labelled set: {len(x["train_labelled"])}')
    print(f'Len unlabelled set: {len(x["train_unlabelled"])}')
    

data/nabirds.py
DTennant/IGCD
import os
import pandas as pd
import numpy as np
from copy import deepcopy
from scipy import io as mat_io

from torchvision.datasets.folder import default_loader
from torch.utils.data import Dataset

from data.data_utils import subsample_instances
from config import car_root

class CarsDataset(Dataset):
    """
        Cars Dataset
    """
    def __init__(self, train=True, limit=0, data_dir=car_root, transform=None):

        metas = os.path.join(data_dir, 'devkit/cars_train_annos.mat') if train else os.path.join(data_dir, 'devkit/cars_test_annos_withlabels.mat')
        data_dir = os.path.join(data_dir, 'cars_train/') if train else os.path.join(data_dir, 'cars_test/')

        self.loader = default_loader
        self.data_dir = data_dir
        self.data = []
        self.target = []
        self.train = train

        self.transform = transform

        if not isinstance(metas, str):
            raise Exception("Train metas must be string location !")
        labels_meta = mat_io.loadmat(metas)

        for idx, img_ in enumerate(labels_meta['annotations'][0]):
            if limit:
                if idx > limit:
                    break

            # self.data.append(img_resized)
            self.data.append(data_dir + img_[5][0])
            # if self.mode == 'train':
            self.target.append(img_[4][0][0])

        self.uq_idxs = np.array(range(len(self)))
        self.target_transform = None

    def __getitem__(self, idx):

        image = self.loader(self.data[idx])
        target = self.target[idx] - 1

        if self.transform is not None:
            image = self.transform(image)

        if self.target_transform is not None:
            target = self.target_transform(target)

        idx = self.uq_idxs[idx]

        return image, target, idx

    def __len__(self):
        return len(self.data)


def subsample_dataset(dataset, idxs):

    dataset.data = np.array(dataset.data)[idxs].tolist()
    dataset.target = np.array(dataset.target)[idxs].tolist()
    dataset.uq_idxs = dataset.uq_idxs[idxs]

    return dataset


def subsample_classes(dataset, include_classes=range(160)):

    include_classes_cars = np.array(include_classes) + 1     # SCars classes are indexed 1 --> 196 instead of 0 --> 195
    cls_idxs = [x for x, t in enumerate(dataset.target) if t in include_classes_cars]

    target_xform_dict = {}
    for i, k in enumerate(include_classes):
        target_xform_dict[k] = i

    dataset = subsample_dataset(dataset, cls_idxs)

    # dataset.target_transform = lambda x: target_xform_dict[x]

    return dataset

def get_train_val_indices(train_dataset, val_split=0.2):

    train_classes = np.unique(train_dataset.target)

    # Get train/test indices
    train_idxs = []
    val_idxs = []
    for cls in train_classes:

        cls_idxs = np.where(train_dataset.target == cls)[0]

        v_ = np.random.choice(cls_idxs, replace=False, size=((int(val_split * len(cls_idxs))),))
        t_ = [x for x in cls_idxs if x not in v_]

        train_idxs.extend(t_)
        val_idxs.extend(v_)

    return train_idxs, val_idxs


def get_scars_datasets(train_transform, test_transform, train_classes=range(160), prop_train_labels=0.8,
                    split_train_val=False, seed=0, args=None):

    np.random.seed(seed)

    # Init entire training set
    whole_training_set = CarsDataset(data_dir=car_root, transform=train_transform, train=True)

    # Get labelled training set which has subsampled classes, then subsample some indices from that
    train_dataset_labelled = subsample_classes(deepcopy(whole_training_set), include_classes=train_classes)
    subsample_indices = subsample_instances(train_dataset_labelled, prop_indices_to_subsample=prop_train_labels)
    train_dataset_labelled = subsample_dataset(train_dataset_labelled, subsample_indices)

    # Split into training and validation sets
    train_idxs, val_idxs = get_train_val_indices(train_dataset_labelled)
    train_dataset_labelled_split = subsample_dataset(deepcopy(train_dataset_labelled), train_idxs)
    val_dataset_labelled_split = subsample_dataset(deepcopy(train_dataset_labelled), val_idxs)
    val_dataset_labelled_split.transform = test_transform

    # Get unlabelled data
    unlabelled_indices = set(whole_training_set.uq_idxs) - set(train_dataset_labelled.uq_idxs)
    train_dataset_unlabelled = subsample_dataset(deepcopy(whole_training_set), np.array(list(unlabelled_indices)))

    # Sample the unlabelled data to only have overlap_cls and unlabelled_classes
    if args.n_overlap_cls != -1:
        cls2sample = list(set(args.overlap_cls).union(set(args.unlabeled_classes)))
        train_dataset_unlabelled = subsample_classes(deepcopy(train_dataset_unlabelled), include_classes=cls2sample)

    # Get test set for all classes
    test_dataset = CarsDataset(data_dir=car_root, transform=test_transform, train=False)

    # Either split train into train and val or use test set as val
    train_dataset_labelled = train_dataset_labelled_split if split_train_val else train_dataset_labelled
    val_dataset_labelled = val_dataset_labelled_split if split_train_val else None

    all_datasets = {
        'train_labelled': train_dataset_labelled,
        'train_unlabelled': train_dataset_unlabelled,
        'val': val_dataset_labelled,
        'test': test_dataset,
    }

    return all_datasets

if __name__ == '__main__':

    x = get_scars_datasets(None, None, train_classes=range(98), prop_train_labels=0.5, split_train_val=False)

    print('Printing lens...')
    for k, v in x.items():
        if v is not None:
            print(f'{k}: {len(v)}')

    print('Printing labelled and unlabelled overlap...')
    print(set.intersection(set(x['train_labelled'].uq_idxs), set(x['train_unlabelled'].uq_idxs)))
    print('Printing total instances in train...')
    print(len(set(x['train_labelled'].uq_idxs)) + len(set(x['train_unlabelled'].uq_idxs)))

    print(f'Num Labelled Classes: {len(set(x["train_labelled"].target))}')
    print(f'Num Unabelled Classes: {len(set(x["train_unlabelled"].target))}')
    print(f'Len labelled set: {len(x["train_labelled"])}')
    print(f'Len unlabelled set: {len(x["train_unlabelled"])}')
data/stanford_cars.py
DTennant/IGCD
import argparse

import math
import numpy as np
import torch
import torch.nn as nn
from torch.optim import SGD, lr_scheduler
from torch.utils.data import DataLoader
from tqdm import tqdm
import wandb
import copy
import pandas as pd

from data.augmentations import get_transform
from data.get_datasets import get_datasets, get_class_splits

from util.general_utils import AverageMeter, init_experiment
from util.cluster_and_log_utils import log_accs_from_preds
from config import exp_root
from model.loss import info_nce_logits, SupConLoss, DistillLoss, ContrastiveLearningViewGenerator, get_params_groups



class iCarlNet(nn.Module):
    def __init__(self, feature_extractor, dino_head, args=None):
        super().__init__()
        self.feature_extractor = feature_extractor
        self.dino_head = dino_head
        self.args = args
        self.n_classes = len(args.classes_in_labelled)
        
        # List containing exemplar_sets
        # each entry is a list of pathes to the images
        """
        [
            [(path1.png, label1), (path2.png, label1), ...],
            [(path1.png, label2), (path2.png, label2), ...],
        ]
        """
        self.exemplar_sets = []

        # Learning method
        self.icarl_cls_loss = nn.CrossEntropyLoss()
        self.dist_loss = nn.BCELoss()

        # Means of exemplars
        self.compute_means = True
        self.exemplar_means = []

    def forward(self, x, return_x=False):
        feature = self.feature_extractor(x)
        output = self.dino_head(feature, return_x=return_x)
        return output

    def incremental_classes(self, n):
        self.n_classes += n
        
    def classify(self, x, transform):
        batch_size = x.shape[0]
        
        if self.compute_means:
            self.args.logger.info("Computing means of exemplars")
            exemplar_means = []
            for exemplar_set in self.exemplar_sets:
                features = []
                for ex, _ in exemplar_set:
                    ex = transform(Image.open(ex)).cuda()
                    ex = ex.unsqueeze(0)
                    feature = self.feature_extractor(ex)
                    features.append(feature)
                features = torch.cat(features, dim=0)
                features = features / features.norm(dim=-1, keepdim=True)
                mean = features.mean(dim=0)
                mean = mean / mean.norm(dim=-1, keepdim=True)
                exemplar_means.append(mean)
            self.exemplar_means = exemplar_means
            self.compute_means = False
            self.args.logger.info("Done computing means of exemplars")

        exemplar_means = self.exemplar_means
        means = torch.stack(exemplar_means, dim=0)
        means = torch.stack([means] * batch_size, dim=0)
        means = torch.transpose(means, 1, 2)
        
        features = self.feature_extractor(x)
        features = features / features.norm(dim=-1, keepdim=True)
        features = features.unsqueeze(2)
        features = features.expand_as(means)
        
        dist = (features - means).pow(2).sum(dim=1).squeeze()
        _, preds = dist.min(dim=1)
        return preds

    @torch.no_grad()
    def construct_exemplar_sets(self, train_loader, n, args):
        # import ipdb; ipdb.set_trace()
        self.args.logger.info("Constructing exemplar sets")
        self.eval()
        preds, labels, mask_labs, paths = [], [], [], []
        feats = []
        for batch_idx, batch in enumerate(tqdm(train_loader)):
            images, class_labels, uq_idxs, path, mask_lab = batch
            paths.append(path)
            images = images[0].cuda(non_blocking=True) # only one view of the image
            features = self.feature_extractor(images)
            features = features / features.norm(dim=-1, keepdim=True)
            _, pred = self.dino_head(features)
            preds.append(pred.argmax(1).cpu().numpy())
            labels.append(class_labels.cpu().numpy())
            mask_labs.append(mask_lab.cpu().numpy())
            feats.append(features.cpu().numpy())
        preds = np.concatenate(preds)
        labels = np.concatenate(labels)
        mask_labs = np.concatenate(mask_labs).astype(bool)
        features = np.concatenate(feats)
        paths = np.concatenate(paths)
        
        labels_used = np.zeros(labels.shape) # when mask_lab is True, use labels, when False, use preds
        labels_used[mask_labs] = labels[mask_labs]
        labels_used[~mask_labs] = preds[~mask_labs]
        labels_used = labels_used.astype(np.int32)
        
        for lbl in np.unique(labels_used):
            feats = features[labels_used == lbl]
            lbl_paths = paths[labels_used == lbl]
            sub_labels_used = labels_used[labels_used == lbl]
            mean_feat = feats.mean(axis=0) # mean of features
            mean_feat = mean_feat / np.linalg.norm(mean_feat) # normalize
            
            exemplar_set = []
            exemplar_features = []
            for k in range(n):
                S = np.sum(exemplar_features, axis=0)
                phi = feats
                mu = mean_feat
                mu_p = 1.0/(k+1) * (phi + S)
                mu_p = mu_p / np.linalg.norm(mu_p)
                i = np.argmin(np.sqrt(np.sum((mu - mu_p) ** 2, axis=1)))

                exemplar_set.append(lbl_paths[i])
                exemplar_features.append(feats[i])

            self.exemplar_sets.append([(item, lbl) for item in exemplar_set])
    
    def reduce_exemplar_sets(self, n):
        for idx, exemplar_set in enumerate(self.exemplar_sets):
            self.exemplar_sets[idx] = exemplar_set[:n]
        
    def combine_dataset_with_exemplars(self, train_dataset, n, args):
        # return train_dataset + self.exemplar_sets
        # import ipdb; ipdb.set_trace()
        self.reduce_exemplar_sets(n)
        file_names, category_ids = [], []
        for exemplar_set in self.exemplar_sets:
            for ex, cat in exemplar_set:
                file_names.append(ex)
                category_ids.append(cat)
        new_df = pd.concat([train_dataset.data, pd.DataFrame({'file_name': file_names, 'category_id': category_ids})])
        new_df = new_df.drop_duplicates(subset=['file_name', ], keep='first')
        new_dataset = copy.deepcopy(train_dataset)
        new_dataset.data = new_df
        return new_dataset


model/icarl.py
DTennant/IGCD
import torch
import torch.nn as nn
import numpy as np
from torch.nn import functional as F

class ContrastiveLearningViewGenerator(object):
    """Take two random crops of one image as the query and key."""

    def __init__(self, base_transform, n_views=2):
        self.base_transform = base_transform
        self.n_views = n_views

    def __call__(self, x):
        if not isinstance(self.base_transform, list):
            return [self.base_transform(x) for i in range(self.n_views)]
        else:
            return [self.base_transform[i](x) for i in range(self.n_views)]

class SupConLoss(torch.nn.Module):
    """Supervised Contrastive Learning: https://arxiv.org/pdf/2004.11362.pdf.
    It also supports the unsupervised contrastive loss in SimCLR
    From: https://github.com/HobbitLong/SupContrast"""
    def __init__(self, temperature=0.07, contrast_mode='all',
                 base_temperature=0.07):
        super(SupConLoss, self).__init__()
        self.temperature = temperature
        self.contrast_mode = contrast_mode
        self.base_temperature = base_temperature

    def forward(self, features, labels=None, mask=None):
        """Compute loss for model. If both `labels` and `mask` are None,
        it degenerates to SimCLR unsupervised loss:
        https://arxiv.org/pdf/2002.05709.pdf
        Args:
            features: hidden vector of shape [bsz, n_views, ...].
            labels: ground truth of shape [bsz].
            mask: contrastive mask of shape [bsz, bsz], mask_{i,j}=1 if sample j
                has the same class as sample i. Can be asymmetric.
        Returns:
            A loss scalar.
        """

        device = (torch.device('cuda')
                  if features.is_cuda
                  else torch.device('cpu'))

        if len(features.shape) < 3:
            raise ValueError('`features` needs to be [bsz, n_views, ...],'
                             'at least 3 dimensions are required')
        if len(features.shape) > 3:
            features = features.view(features.shape[0], features.shape[1], -1)

        batch_size = features.shape[0]
        if labels is not None and mask is not None:
            raise ValueError('Cannot define both `labels` and `mask`')
        elif labels is None and mask is None:
            mask = torch.eye(batch_size, dtype=torch.float32).to(device)
        elif labels is not None:
            labels = labels.contiguous().view(-1, 1)
            if labels.shape[0] != batch_size:
                raise ValueError('Num of labels does not match num of features')
            mask = torch.eq(labels, labels.T).float().to(device)
        else:
            mask = mask.float().to(device)

        contrast_count = features.shape[1]
        contrast_feature = torch.cat(torch.unbind(features, dim=1), dim=0)
        if self.contrast_mode == 'one':
            anchor_feature = features[:, 0]
            anchor_count = 1
        elif self.contrast_mode == 'all':
            anchor_feature = contrast_feature
            anchor_count = contrast_count
        else:
            raise ValueError('Unknown mode: {}'.format(self.contrast_mode))

        # compute logits
        anchor_dot_contrast = torch.div(
            torch.matmul(anchor_feature, contrast_feature.T),
            self.temperature)

        # for numerical stability
        logits_max, _ = torch.max(anchor_dot_contrast, dim=1, keepdim=True)
        logits = anchor_dot_contrast - logits_max.detach()

        # tile mask
        mask = mask.repeat(anchor_count, contrast_count)
        # mask-out self-contrast cases
        logits_mask = torch.scatter(
            torch.ones_like(mask),
            1,
            torch.arange(batch_size * anchor_count).view(-1, 1).to(device),
            0
        )
        mask = mask * logits_mask

        # compute log_prob
        exp_logits = torch.exp(logits) * logits_mask
        log_prob = logits - torch.log(exp_logits.sum(1, keepdim=True))

        # compute mean of log-likelihood over positive
        mean_log_prob_pos = (mask * log_prob).sum(1) / mask.sum(1)

        # loss
        loss = - (self.temperature / self.base_temperature) * mean_log_prob_pos
        loss = loss.view(anchor_count, batch_size).mean()

        return loss



def info_nce_logits(features, n_views=2, temperature=1.0, device='cuda'):

    b_ = 0.5 * int(features.size(0))

    labels = torch.cat([torch.arange(b_) for i in range(n_views)], dim=0)
    labels = (labels.unsqueeze(0) == labels.unsqueeze(1)).float()
    labels = labels.to(device)

    features = F.normalize(features, dim=1)

    similarity_matrix = torch.matmul(features, features.T)

    # discard the main diagonal from both: labels and similarities matrix
    mask = torch.eye(labels.shape[0], dtype=torch.bool).to(device)
    labels = labels[~mask].view(labels.shape[0], -1)
    similarity_matrix = similarity_matrix[~mask].view(similarity_matrix.shape[0], -1)

    # select and combine multiple positives
    positives = similarity_matrix[labels.bool()].view(labels.shape[0], -1)

    # select only the negatives the negatives
    negatives = similarity_matrix[~labels.bool()].view(similarity_matrix.shape[0], -1)

    logits = torch.cat([positives, negatives], dim=1)
    labels = torch.zeros(logits.shape[0], dtype=torch.long).to(device)

    logits = logits / temperature
    return logits, labels


def get_params_groups(model):
    regularized = []
    not_regularized = []
    for name, param in model.named_parameters():
        if not param.requires_grad:
            continue
        # we do not regularize biases nor Norm parameters
        if name.endswith(".bias") or len(param.shape) == 1:
            not_regularized.append(param)
        else:
            regularized.append(param)
    return [{'params': regularized}, {'params': not_regularized, 'weight_decay': 0.}]


class DistillLoss(nn.Module):
    def __init__(self, warmup_teacher_temp_epochs, nepochs, 
                 ncrops=2, warmup_teacher_temp=0.07, teacher_temp=0.04,
                 student_temp=0.1):
        super().__init__()
        self.student_temp = student_temp
        self.ncrops = ncrops
        self.teacher_temp_schedule = np.concatenate((
            np.linspace(warmup_teacher_temp,
                        teacher_temp, warmup_teacher_temp_epochs),
            np.ones(nepochs - warmup_teacher_temp_epochs) * teacher_temp
        ))

    def forward(self, student_output, teacher_output, epoch):
        """
        Cross-entropy between softmax outputs of the teacher and student networks.
        """
        student_out = student_output / self.student_temp
        student_out = student_out.chunk(self.ncrops)

        # teacher centering and sharpening
        temp = self.teacher_temp_schedule[epoch]
        teacher_out = F.softmax(teacher_output / temp, dim=-1)
        teacher_out = teacher_out.detach().chunk(2)

        total_loss = 0
        n_loss_terms = 0
        for iq, q in enumerate(teacher_out):
            for v in range(len(student_out)):
                if v == iq:
                    # we skip cases where student and teacher operate on the same view
                    continue
                loss = torch.sum(-q * F.log_softmax(student_out[v], dim=-1), dim=-1)
                total_loss += loss.mean()
                n_loss_terms += 1
        total_loss /= n_loss_terms
        return total_loss

model/loss.py
DTennant/IGCD
import argparse
import torch
torch.multiprocessing.set_sharing_strategy('file_system')

import math
import numpy as np
import torch
import torch.nn as nn
from torch.optim import SGD, lr_scheduler
from torch.utils.data import DataLoader
from tqdm import tqdm
import wandb

from data.augmentations import get_transform
from data.get_datasets import get_datasets, get_class_splits

from util.general_utils import AverageMeter, init_experiment
from util.cluster_and_log_utils import log_accs_from_preds
from config import exp_root
from model.loss import info_nce_logits, SupConLoss, DistillLoss, ContrastiveLearningViewGenerator, get_params_groups


class DINOHead(nn.Module):
    def __init__(self, in_dim, out_dim, use_bn=False, norm_last_layer=True, 
                 nlayers=3, hidden_dim=2048, bottleneck_dim=256):
        super().__init__()
        nlayers = max(nlayers, 1)
        if nlayers == 1:
            self.mlp = nn.Linear(in_dim, bottleneck_dim)
        elif nlayers != 0:
            layers = [nn.Linear(in_dim, hidden_dim)]
            if use_bn:
                layers.append(nn.BatchNorm1d(hidden_dim))
            layers.append(nn.GELU())
            for _ in range(nlayers - 2):
                layers.append(nn.Linear(hidden_dim, hidden_dim))
                if use_bn:
                    layers.append(nn.BatchNorm1d(hidden_dim))
                layers.append(nn.GELU())
            layers.append(nn.Linear(hidden_dim, bottleneck_dim))
            self.mlp = nn.Sequential(*layers)
        self.apply(self._init_weights)
        self.last_layer = nn.utils.weight_norm(nn.Linear(in_dim, out_dim, bias=False))
        self.last_layer.weight_g.data.fill_(1)
        if norm_last_layer:
            self.last_layer.weight_g.requires_grad = False

    def _init_weights(self, m):
        if isinstance(m, nn.Linear):
            torch.nn.init.trunc_normal_(m.weight, std=.02)
            if isinstance(m, nn.Linear) and m.bias is not None:
                nn.init.constant_(m.bias, 0)

    def forward(self, x):
        x_proj = self.mlp(x)
        x = nn.functional.normalize(x, dim=-1, p=2)
        # x = x.detach()
        logits = self.last_layer(x)
        return x_proj, logits


def train(student, train_loader, test_loader, unlabelled_train_loader, args):
    params_groups = get_params_groups(student)
    optimizer = SGD(params_groups, lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay)
    fp16_scaler = None
    if args.fp16:
        fp16_scaler = torch.cuda.amp.GradScaler()

    exp_lr_scheduler = lr_scheduler.CosineAnnealingLR(
            optimizer,
            T_max=args.epochs,
            eta_min=args.lr * 1e-3,
        )

    if args.resume:
        load_dict = torch.load(args.resume, map_location='cpu')
        model_dict = load_dict['model']
        optimizer_dict = load_dict['optimizer']
        lr_scheduler_dict = load_dict['lr_scheduler']
        args.start_epoch = load_dict['epoch']
        
        student.load_state_dict(model_dict)
        optimizer.load_state_dict(optimizer_dict)
        exp_lr_scheduler.load_state_dict(lr_scheduler_dict)
        
        args.logger.info(f"Resumed from {args.resume}")

    cluster_criterion = DistillLoss(
                        args.warmup_teacher_temp_epochs,
                        args.epochs,
                        args.n_views,
                        args.warmup_teacher_temp,
                        args.teacher_temp,
                    )

    # inductive
    best_test_acc_ubl = 0
    # transductive
    best_train_acc_lab = 0
    best_train_acc_ubl = 0 
    best_train_acc_all = 0

    for epoch in range(args.start_epoch, args.epochs):
        loss_record = AverageMeter()

        student.train()
        for batch_idx, batch in enumerate(train_loader):
            images, class_labels, uq_idxs, mask_lab = batch
            # mask_lab = mask_lab[:, 0]

            class_labels, mask_lab = class_labels.cuda(non_blocking=True), mask_lab.cuda(non_blocking=True).bool()
            images = torch.cat(images, dim=0).cuda(non_blocking=True)

            with torch.cuda.amp.autocast(fp16_scaler is not None):
                student_proj, student_out = student(images)
                teacher_out = student_out.detach()

                # clustering, sup
                sup_logits = torch.cat([f[mask_lab] for f in (student_out / 0.1).chunk(2)], dim=0)
                sup_labels = torch.cat([class_labels[mask_lab] for _ in range(2)], dim=0)
                cls_loss = nn.CrossEntropyLoss()(sup_logits, sup_labels)

                # clustering, unsup
                cluster_loss = cluster_criterion(student_out, teacher_out, epoch)
                avg_probs = (student_out / 0.1).softmax(dim=1).mean(dim=0)
                me_max_loss = - torch.sum(torch.log(avg_probs**(-avg_probs))) + math.log(float(len(avg_probs)))
                cluster_loss += args.memax_weight * me_max_loss

                # represent learning, unsup
                contrastive_logits, contrastive_labels = info_nce_logits(features=student_proj)
                contrastive_loss = torch.nn.CrossEntropyLoss()(contrastive_logits, contrastive_labels)

                # representation learning, sup
                student_proj = torch.cat([f[mask_lab].unsqueeze(1) for f in student_proj.chunk(2)], dim=1)
                student_proj = torch.nn.functional.normalize(student_proj, dim=-1)
                sup_con_labels = class_labels[mask_lab]
                sup_con_loss = SupConLoss()(student_proj, labels=sup_con_labels)

                pstr = ''
                pstr += f'cls_loss: {cls_loss.item():.4f} '
                pstr += f'cluster_loss: {cluster_loss.item():.4f} '
                pstr += f'sup_con_loss: {sup_con_loss.item():.4f} '
                pstr += f'contrastive_loss: {contrastive_loss.item():.4f} '

                loss = 0
                loss += (1 - args.sup_weight) * cluster_loss + args.sup_weight * cls_loss
                loss += (1 - args.sup_weight) * contrastive_loss + args.sup_weight * sup_con_loss
                
            # Train acc
            loss_record.update(loss.item(), class_labels.size(0))
            optimizer.zero_grad()
            if fp16_scaler is None:
                loss.backward()
                optimizer.step()
            else:
                fp16_scaler.scale(loss).backward()
                fp16_scaler.step(optimizer)
                fp16_scaler.update()

            if batch_idx % args.print_freq == 0:
                args.logger.info('Epoch: [{}][{}/{}]\t loss {:.5f}\t {}'
                            .format(epoch, batch_idx, len(train_loader), loss.item(), pstr))
                
            wandb.log({
                'cls_loss': cls_loss.item(),
                'cluster_loss': cluster_loss.item(),
                'sup_con_loss': sup_con_loss.item(),
                'contrastive_loss': contrastive_loss.item(),
                'loss': loss.item(),
                'lr': optimizer.param_groups[0]['lr'],
            })

        args.logger.info('Train Epoch: {} Avg Loss: {:.4f} '.format(epoch, loss_record.avg))

        args.logger.info('Testing on unlabelled examples in the training data...')
        output = test(student, unlabelled_train_loader, epoch=epoch, save_name='Train ACC Unlabelled', args=args)
        all_acc, old_acc, new_acc = output['All'], output['Old'], output['New']
        args.logger.info('Testing on disjoint test set...')
        output_test = test(student, test_loader, epoch=epoch, save_name='Test ACC', args=args)
        all_acc_test, old_acc_test, new_acc_test = output_test['All'], output_test['Old'], output_test['New']
        
        wandb.log({k + '_train': v for k, v in output.items()})
        wandb.log({k + '_test': v for k, v in output_test.items()})


        args.logger.info('Train Accuracies: All {:.4f} | Old {:.4f} | New {:.4f}'.format(all_acc, old_acc, new_acc))
        args.logger.info('Test Accuracies: All {:.4f} | Old {:.4f} | New {:.4f}'.format(all_acc_test, old_acc_test, new_acc_test))

        # Step schedule
        exp_lr_scheduler.step()

        save_dict = {
            'model': student.state_dict(),
            'optimizer': optimizer.state_dict(),
            'lr_scheduler': exp_lr_scheduler.state_dict(),
            'epoch': epoch,
        }

        torch.save(save_dict, args.model_path)
        args.logger.info("model saved to {}.".format(args.model_path))

        if new_acc_test > best_test_acc_ubl:

            args.logger.info(f'Best ACC on old Classes on disjoint test set: {old_acc_test:.4f}...')
            args.logger.info('Best Train Accuracies: All {:.4f} | Old {:.4f} | New {:.4f}'.format(all_acc, old_acc, new_acc))

            torch.save(save_dict, args.model_path[:-3] + f'_best.pt')
            args.logger.info("model saved to {}.".format(args.model_path[:-3] + f'_best.pt'))

            # inductive
            best_output_test = output_test
            best_test_acc_ubl = new_acc_test
            # transductive            
            best_output = output
            best_train_acc_lab = old_acc
            best_train_acc_ubl = new_acc
            best_train_acc_all = all_acc

        args.logger.info(f'Exp Name: {args.exp_name} Exp id: {args.exp_id}')
        best_pstr = f'Metrics with best model on test set: '
        output_dict = best_output_test if args.eval_setting == 'inductive' else best_output
        for k, v in output_dict.items():
            best_pstr += f'{k}: {v:.4f} '
            wandb.log({f'best_{k}': v})
        args.logger.info(best_pstr)


def test(model, test_loader, epoch, save_name, args):

    model.eval()

    preds, targets = [], []
    mask = np.array([])
    for batch_idx, (images, label, _) in enumerate(tqdm(test_loader)):
        images = images.cuda(non_blocking=True)
        with torch.no_grad():
            _, logits = model(images)
            preds.append(logits.argmax(1).cpu().numpy())
            targets.append(label.cpu().numpy())
            mask = np.append(mask, np.array([True if x.item() in range(len(args.train_classes)) else False for x in label]))

    preds = np.concatenate(preds)
    targets = np.concatenate(targets)
    return log_accs_from_preds(y_true=targets, y_pred=preds, mask=mask,
                                                    T=epoch, eval_funcs=args.eval_funcs, save_name=save_name,
                                                    args=args)



if __name__ == "__main__":

    parser = argparse.ArgumentParser(description='cluster', formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument('--batch_size', default=128, type=int)
    parser.add_argument('--num_workers', default=8, type=int)
    parser.add_argument('--eval_funcs', nargs='+', help='Which eval functions to use', default=['v2', ])

    parser.add_argument('--warmup_model_dir', type=str, default=None)
    parser.add_argument('--resume', type=str, default=None)
    
    parser.add_argument('--n_overlap_cls', type=int, default=-1, help='how many classes overlaps, -1 means all labelled cls are in unlabelled')
    parser.add_argument('--dataset_name', type=str, default='scars', help='options: cifar10, cifar100, imagenet_100, cub, scars, fgvc_aricraft, herbarium_19')
    parser.add_argument('--prop_train_labels', type=float, default=0.5)
    parser.add_argument('--use_ssb_splits', action='store_true', default=True)
    parser.add_argument('--eval_setting', type=str, default='transductive', choices=['transductive', 'inductive'], help='options: inductive, transductive')

    parser.add_argument('--grad_from_block', type=int, default=11)
    parser.add_argument('--lr', type=float, default=0.1)
    parser.add_argument('--gamma', type=float, default=0.1)
    parser.add_argument('--momentum', type=float, default=0.9)
    parser.add_argument('--weight_decay', type=float, default=1e-4)
    parser.add_argument('--start_epoch', default=0, type=int)
    parser.add_argument('--epochs', default=200, type=int)
    parser.add_argument('--exp_root', type=str, default=exp_root)
    parser.add_argument('--transform', type=str, default='imagenet')
    parser.add_argument('--sup_weight', type=float, default=0.35)
    parser.add_argument('--n_views', default=2, type=int)
    
    parser.add_argument('--init_estimate_num_cls', type=int, help='the initial estimate of the number of classes, None if use GT', default=None)
    parser.add_argument('--memax_weight', type=float, default=2)
    parser.add_argument('--warmup_teacher_temp', default=0.07, type=float, help='Initial value for the teacher temperature.')
    parser.add_argument('--teacher_temp', default=0.04, type=float, help='Final value (after linear warmup)of the teacher temperature.')
    parser.add_argument('--warmup_teacher_temp_epochs', default=30, type=int, help='Number of warmup epochs for the teacher temperature.')

    parser.add_argument('--imbalance_factor', type=float, default=0.02)

    parser.add_argument('--fp16', action='store_true', default=False)
    parser.add_argument('--print_freq', default=10, type=int)
    parser.add_argument('--runner_name', default='simgcd', type=str, )
    parser.add_argument('--exp_name', default=None, type=str)
    parser.add_argument('--exp_id', default=None, type=str)
    parser.add_argument('--tags', default=None, type=str, help='- separated tags for the experiment; e.g. baseline-ssb')

    # ----------------------
    # INIT
    # ----------------------
    args = parser.parse_args()
    device = torch.device('cuda:0')
    args = get_class_splits(args)

    args.num_labeled_classes = len(args.train_classes)
    args.num_unlabeled_classes = len(args.unlabeled_classes)

    init_experiment(args, runner_name=[args.runner_name], exp_id=args.exp_id)
    args.logger.info(f'Using evaluation function {args.eval_funcs[0]} to print results')
    
    torch.backends.cudnn.benchmark = True

    # ----------------------
    # BASE MODEL
    # ----------------------
    args.interpolation = 3
    args.crop_pct = 0.875

    backbone = torch.hub.load('facebookresearch/dino:main', 'dino_vitb16')

    if args.warmup_model_dir is not None:
        args.logger.info(f'Loading weights from {args.warmup_model_dir}')
        backbone.load_state_dict(torch.load(args.warmup_model_dir, map_location='cpu'))
    
    # NOTE: Hardcoded image size as we do not finetune the entire ViT model
    args.image_size = 224
    args.feat_dim = 768
    args.num_mlp_layers = 3
    args.mlp_out_dim = args.num_labeled_classes + args.num_unlabeled_classes
    if args.init_estimate_num_cls is not None:
        args.mlp_out_dim = args.init_estimate_num_cls

    # ----------------------
    # HOW MUCH OF BASE MODEL TO FINETUNE
    # ----------------------
    for m in backbone.parameters():
        m.requires_grad = False

    # Only finetune layers from block 'args.grad_from_block' onwards
    for name, m in backbone.named_parameters():
        if 'block' in name:
            block_num = int(name.split('.')[1])
            if block_num >= args.grad_from_block:
                m.requires_grad = True

    
    args.logger.info('model build')
    
    # --------------------
    # CONTRASTIVE TRANSFORM
    # --------------------
    train_transform, test_transform = get_transform(args.transform, image_size=args.image_size, args=args)
    train_transform = ContrastiveLearningViewGenerator(base_transform=train_transform, n_views=args.n_views)
    # --------------------
    # DATASETS
    # --------------------
    train_dataset, test_dataset, unlabelled_train_examples_test, datasets = get_datasets(args.dataset_name,
                                                                                         train_transform,
                                                                                         test_transform,
                                                                                         args)
    args.target_transform = train_dataset.labelled_dataset.target_transform

    # --------------------
    # SAMPLER
    # Sampler which balances labelled and unlabelled examples in each batch
    # --------------------
    label_len = len(train_dataset.labelled_dataset)
    unlabelled_len = len(train_dataset.unlabelled_dataset)
    args.logger.info(f'label length: {label_len}, unlabel length: {unlabelled_len}')
    sample_weights = [1 if i < label_len else label_len / unlabelled_len for i in range(len(train_dataset))]
    sample_weights = torch.DoubleTensor(sample_weights)
    sampler = torch.utils.data.WeightedRandomSampler(sample_weights, num_samples=len(train_dataset))

    # --------------------
    # DATALOADERS
    # --------------------
    train_loader = DataLoader(train_dataset, num_workers=args.num_workers, batch_size=args.batch_size, shuffle=False,
                              sampler=sampler, drop_last=True, pin_memory=True)
    test_loader_unlabelled = DataLoader(unlabelled_train_examples_test, num_workers=args.num_workers,
                                        batch_size=256, shuffle=False, pin_memory=False)
    test_loader_labelled = DataLoader(test_dataset, num_workers=args.num_workers,
                                      batch_size=256, shuffle=False, pin_memory=False)

    # ----------------------
    # PROJECTION HEAD
    # ----------------------
    projector = DINOHead(in_dim=args.feat_dim, out_dim=args.mlp_out_dim, nlayers=args.num_mlp_layers)
    model = nn.Sequential(backbone, projector).to(device)

    # ----------------------
    # TRAIN
    # ----------------------
    wandb.init(project='SimGCD', name=args.exp_name, tags=args.tags.split('-') if args.tags is not None else None, resume=True if args.resume is not None else None)
    wandb.config.update(args)

    train(model, train_loader, test_loader_labelled, test_loader_unlabelled, args)

model/simgcd.py
DTennant/IGCD
import argparse

import math
import numpy as np
import torch
import torch.nn as nn
from torch.optim import SGD, lr_scheduler
from torch.utils.data import DataLoader
from tqdm import tqdm
import wandb

from data.augmentations import get_transform
from data.get_datasets import get_datasets, get_class_splits

from util.general_utils import AverageMeter, init_experiment
from util.cluster_and_log_utils import log_accs_from_preds
from config import exp_root
from model.loss import info_nce_logits, SupConLoss, DistillLoss, ContrastiveLearningViewGenerator, get_params_groups

from model.icarl import iCarlNet

class DINOHead(nn.Module):
    def __init__(self, in_dim, out_dim, use_bn=False, norm_last_layer=True, 
                 nlayers=3, hidden_dim=2048, bottleneck_dim=256):
        super().__init__()
        nlayers = max(nlayers, 1)
        if nlayers == 1:
            self.mlp = nn.Linear(in_dim, bottleneck_dim)
        elif nlayers != 0:
            layers = [nn.Linear(in_dim, hidden_dim)]
            if use_bn:
                layers.append(nn.BatchNorm1d(hidden_dim))
            layers.append(nn.GELU())
            for _ in range(nlayers - 2):
                layers.append(nn.Linear(hidden_dim, hidden_dim))
                if use_bn:
                    layers.append(nn.BatchNorm1d(hidden_dim))
                layers.append(nn.GELU())
            layers.append(nn.Linear(hidden_dim, bottleneck_dim))
            self.mlp = nn.Sequential(*layers)
        self.apply(self._init_weights)
        self.last_layer = nn.utils.weight_norm(nn.Linear(in_dim, out_dim, bias=False))
        self.last_layer.weight_g.data.fill_(1)
        self.in_dim = in_dim
        self.out_dim = out_dim
        if norm_last_layer:
            self.last_layer.weight_g.requires_grad = False

    def _init_weights(self, m):
        if isinstance(m, nn.Linear):
            torch.nn.init.trunc_normal_(m.weight, std=.02)
            if isinstance(m, nn.Linear) and m.bias is not None:
                nn.init.constant_(m.bias, 0)

    def forward(self, x, return_x=False):
        x_proj = self.mlp(x)
        x = nn.functional.normalize(x, dim=-1, p=2)
        # x = x.detach()
        logits = self.last_layer(x)
        if return_x:
            return x_proj, logits, x
        return x_proj, logits
    


def train(student, train_loaders, test_loaders, unlabelled_train_loaders, args):
    params_groups = get_params_groups(student)
    optimizer = SGD(params_groups, lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay)
    fp16_scaler = None
    if args.fp16:
        fp16_scaler = torch.cuda.amp.GradScaler()

    exp_lr_scheduler = lr_scheduler.CosineAnnealingLR(
            optimizer,
            T_max=len(train_loaders[-1]) * (args.epochs - args.start_epoch),
            eta_min=args.lr * 1e-3,
        )

    if args.resume:
        load_dict = torch.load(args.resume, map_location='cpu')
        model_dict = load_dict['model']
        optimizer_dict = load_dict['optimizer']
        lr_scheduler_dict = load_dict['lr_scheduler']
        args.start_epoch = load_dict['epoch']
        
        student.load_state_dict(model_dict)
        optimizer.load_state_dict(optimizer_dict)
        exp_lr_scheduler.load_state_dict(lr_scheduler_dict)
        
        args.logger.info(f"Resumed from {args.resume}")

    cluster_criterion = DistillLoss(
                        args.warmup_teacher_temp_epochs,
                        args.epochs,
                        args.n_views,
                        args.warmup_teacher_temp,
                        args.teacher_temp,
                    )

    # inductive
    best_test_acc_ubl = 0
    # transductive
    best_train_acc_lab = 0
    best_train_acc_ubl = 0 
    best_train_acc_all = 0

    for epoch in range(args.start_epoch, args.epochs):
        loss_record = AverageMeter()

        student.train()
        for batch_idx, batch in enumerate(train_loaders[-1]):
            images, class_labels, uq_idxs, _, mask_lab = batch
            # mask_lab = mask_lab[:, 0]

            class_labels, mask_lab = class_labels.cuda(non_blocking=True), mask_lab.cuda(non_blocking=True).bool()
            images = torch.cat(images, dim=0).cuda(non_blocking=True)

            with torch.cuda.amp.autocast(fp16_scaler is not None):
                student_proj, student_out = student(images)
                teacher_out = student_out.detach()

                # clustering, sup
                sup_logits = torch.cat([f[mask_lab] for f in (student_out / 0.1).chunk(2)], dim=0)
                sup_labels = torch.cat([class_labels[mask_lab] for _ in range(2)], dim=0)
                cls_loss = nn.CrossEntropyLoss()(sup_logits, sup_labels)

                # clustering, unsup
                cluster_loss = cluster_criterion(student_out, teacher_out, epoch)
                avg_probs = (student_out / 0.1).softmax(dim=1).mean(dim=0)
                me_max_loss = - torch.sum(torch.log(avg_probs**(-avg_probs))) + math.log(float(len(avg_probs)))
                cluster_loss += args.memax_weight * me_max_loss

                # represent learning, unsup
                contrastive_logits, contrastive_labels = info_nce_logits(features=student_proj)
                contrastive_loss = torch.nn.CrossEntropyLoss()(contrastive_logits, contrastive_labels)

                # representation learning, sup
                student_proj = torch.cat([f[mask_lab].unsqueeze(1) for f in student_proj.chunk(2)], dim=1)
                student_proj = torch.nn.functional.normalize(student_proj, dim=-1)
                sup_con_labels = class_labels[mask_lab]
                sup_con_loss = SupConLoss()(student_proj, labels=sup_con_labels)

                pstr = ''
                pstr += f'cls_loss: {cls_loss.item():.4f} '
                pstr += f'cluster_loss: {cluster_loss.item():.4f} '
                pstr += f'sup_con_loss: {sup_con_loss.item():.4f} '
                pstr += f'contrastive_loss: {contrastive_loss.item():.4f} '

                loss = 0
                loss += (1 - args.sup_weight) * cluster_loss + args.sup_weight * cls_loss
                loss += (1 - args.sup_weight) * contrastive_loss + args.sup_weight * sup_con_loss
                
            # Train acc
            loss_record.update(loss.item(), class_labels.size(0))
            optimizer.zero_grad()
            if fp16_scaler is None:
                loss.backward()
                optimizer.step()
            else:
                fp16_scaler.scale(loss).backward()
                fp16_scaler.step(optimizer)
                fp16_scaler.update()
            # Step schedule
            exp_lr_scheduler.step()

            if batch_idx % args.print_freq == 0:
                args.logger.info('Epoch: [{}][{}/{}]\t loss {:.5f}\t {}'
                            .format(epoch, batch_idx, len(train_loaders[-1]), loss.item(), pstr))
                
            wandb.log({
                'cls_loss': cls_loss.item(),
                'cluster_loss': cluster_loss.item(),
                'sup_con_loss': sup_con_loss.item(),
                'contrastive_loss': contrastive_loss.item(),
                'loss': loss.item(),
                'lr': optimizer.param_groups[0]['lr'],
            })

        args.logger.info('Train Epoch: {} Avg Loss: {:.4f} '.format(epoch, loss_record.avg))

        save_dict = {
            'model': student.state_dict(),
            'optimizer': optimizer.state_dict(),
            'lr_scheduler': exp_lr_scheduler.state_dict(),
            'epoch': epoch,
        }
        torch.save(save_dict, args.model_path[:-3] + f'_stage_{len(train_loaders)}.pth')
        args.logger.info("model saved to {}.".format(args.model_path[:-3] + f'_stage_{len(train_loaders)}.pth'))

        # args.logger.info('Testing on unlabelled examples in the training data...')
        # outputs = test_multi_test_loader(student, unlabelled_train_loaders, epoch=epoch, save_name='Train ACC Unlabelled', args=args)

        # args.logger.info('Testing on disjoint test set...')
        # outputs_test = test_multi_test_loader(student, test_loaders, epoch=epoch, save_name='Test ACC', args=args)

        for stage_idx, (unlabelled_train_loader, test_loader, classes_in_labelled, classes_in_unlabelled, ) in enumerate(zip(unlabelled_train_loaders, test_loaders, 
                                                                                                    args.classes_in_labelled_list, args.classes_in_unlabelled_list, )):
            # NOTE: here we need make sure these classes are in the target transform
            args.classes_in_labelled, args.classes_in_unlabelled = classes_in_labelled, classes_in_unlabelled
            stage_idx += 1

            args.logger.info('Testing on unlabelled examples in the training data...')
            output = test(student, unlabelled_train_loader, epoch=epoch, save_name=f'Train ACC Unlabelled stage {stage_idx}', stage_idx=stage_idx-1, args=args) 

            args.logger.info('Testing on disjoint test set...')
            output_test = test(student, test_loader, epoch=epoch, save_name=f'Test ACC stage {stage_idx}', stage_idx=stage_idx-1, args=args)

            all_acc, old_acc, new_acc = output['All'], output['Old'], output['New']
            all_acc_test, old_acc_test, new_acc_test = output_test['All'], output_test['Old'], output_test['New']
            
            wandb.log({f'stage_{stage_idx}_{len(train_loaders)}_' + k + '_train': v for k, v in output.items()})
            wandb.log({f'stage_{stage_idx}_{len(train_loaders)}_' + k + '_test': v for k, v in output_test.items()})

            args.logger.info('Train Accuracies stage {}: All {:.4f} | Old {:.4f} | New {:.4f}'.format(stage_idx, all_acc, old_acc, new_acc))
            args.logger.info('Test Accuracies stage {}: All {:.4f} | Old {:.4f} | New {:.4f}'.format(stage_idx, all_acc_test, old_acc_test, new_acc_test))

            if stage_idx == len(train_loaders):
                if new_acc_test > best_test_acc_ubl:

                    args.logger.info(f'Best ACC on old Classes on disjoint test set: {old_acc_test:.4f}...')
                    args.logger.info('Best Train Accuracies: All {:.4f} | Old {:.4f} | New {:.4f}'.format(all_acc, old_acc, new_acc))

                    torch.save(save_dict, args.model_path[:-3] + f'_{stage_idx}_{len(train_loaders)}_best.pt')
                    args.logger.info("model saved to {}.".format(args.model_path[:-3] + f'_{stage_idx}_best.pt'))

                    # inductive
                    best_output_test = output_test
                    best_test_acc_ubl = new_acc_test
                    # transductive            
                    best_output = output
                    best_train_acc_lab = old_acc
                    best_train_acc_ubl = new_acc
                    best_train_acc_all = all_acc

                args.logger.info(f'Exp Name: {args.exp_name} Exp id: {args.exp_id}')
                best_pstr = f'Metrics with best model on test set: '
                output_dict = best_output_test if args.eval_setting == 'inductive' else best_output
                for k, v in output_dict.items():
                    best_pstr += f'{k}: {v:.4f} '
                    wandb.log({f'stage_{stage_idx}_{len(train_loaders)}_best_{k}': v})
                args.logger.info(best_pstr)

def test(model, test_loader, epoch, save_name, stage_idx, args):
    model.eval()
    cur_stage = args.stage - 1

    preds, targets = [], []
    mask = np.array([])
    for batch_idx, (images, label, _, _) in enumerate(tqdm(test_loader)):
        images = images.cuda(non_blocking=True)
        with torch.no_grad():
            _, logits = model(images)
            preds.append(logits.argmax(1).cpu().numpy())
            targets.append(label.cpu().numpy())
            mask = np.append(mask, np.array([True if x.item() in range(len(args.train_classes)) else False for x in label]))

    preds = np.concatenate(preds)
    targets = np.concatenate(targets)
    return log_accs_from_preds(y_true=targets, y_pred=preds, mask=mask,
                                                    T=epoch, eval_funcs=args.eval_funcs, save_name=save_name,
                                                    args=args)




if __name__ == "__main__":

    parser = argparse.ArgumentParser(description='cluster', formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument('--batch_size', default=128, type=int)
    parser.add_argument('--num_workers', default=8, type=int)
    parser.add_argument('--eval_funcs', nargs='+', help='Which eval functions to use', default=['v2', ])

    parser.add_argument('--warmup_model_dir', type=str, default=None)
    parser.add_argument('--resume', type=str, default=None)
    
    parser.add_argument('--n_overlap_cls', type=int, default=-1, help='how many classes overlaps, -1 means all labelled cls are in unlabelled')
    parser.add_argument('--dataset_name', type=str, default='scars', help='options: cifar10, cifar100, imagenet_100, cub, scars, fgvc_aricraft, herbarium_19')
    parser.add_argument('--prop_train_labels', type=float, default=0.5)
    parser.add_argument('--use_ssb_splits', action='store_true', default=True)
    parser.add_argument('--eval_setting', type=str, default='transductive', choices=['transductive', 'inductive'], help='options: inductive, transductive')

    parser.add_argument('--grad_from_block', type=int, default=11)
    parser.add_argument('--lr', type=float, default=0.1)
    parser.add_argument('--gamma', type=float, default=0.1)
    parser.add_argument('--momentum', type=float, default=0.9)
    parser.add_argument('--weight_decay', type=float, default=1e-4)
    parser.add_argument('--start_epoch', default=0, type=int)
    parser.add_argument('--epochs', default=200, type=int)
    parser.add_argument('--exp_root', type=str, default=exp_root)
    parser.add_argument('--transform', type=str, default='imagenet')
    parser.add_argument('--sup_weight', type=float, default=0.35)
    parser.add_argument('--n_views', default=2, type=int)
    
    parser.add_argument('--init_estimate_num_cls', type=int, help='the initial estimate of the number of classes, None if use GT', default=None)
    parser.add_argument('--memax_weight', type=float, default=2)
    parser.add_argument('--warmup_teacher_temp', default=0.07, type=float, help='Initial value for the teacher temperature.')
    parser.add_argument('--teacher_temp', default=0.04, type=float, help='Final value (after linear warmup)of the teacher temperature.')
    parser.add_argument('--warmup_teacher_temp_epochs', default=30, type=int, help='Number of warmup epochs for the teacher temperature.')

    parser.add_argument('--fp16', action='store_true', default=False)

    parser.add_argument('--use_small_set', action='store_true', default=False)
    parser.add_argument('--split_crit', type=str, )

    parser.add_argument('--exemplar_number', type=int, default=5)
    parser.add_argument('--return_path', action='store_true', default=False)
    
    parser.add_argument('--print_freq', default=10, type=int)
    parser.add_argument('--runner_name', default='simgcd', type=str, )
    parser.add_argument('--entity', default='tennant', type=str, help='entity for wandb')
    parser.add_argument('--exp_name', default=None, type=str)
    parser.add_argument('--exp_id', default=None, type=str)
    parser.add_argument('--tags', default=None, type=str, help='- separated tags for the experiment; e.g. baseline-ssb')

    # ----------------------
    # INIT
    # ----------------------
    args = parser.parse_args()
    device = torch.device('cuda:0')
    args = get_class_splits(args)

    args.num_labeled_classes = len(args.train_classes)
    args.num_unlabeled_classes = len(args.unlabeled_classes)

    init_experiment(args, runner_name=[args.runner_name], exp_id=args.exp_id)
    args.logger.info(f'Using evaluation function {args.eval_funcs[0]} to print results')
    
    torch.backends.cudnn.benchmark = True

    if args.split_crit in ['year', 'loc_year']:
        args.max_stage = 4
    elif args.split_crit in ['location']:
        args.max_stage = 3
        

    # ----------------------
    # BASE MODEL
    # ----------------------
    args.interpolation = 3
    args.crop_pct = 0.875

    backbone = torch.hub.load('facebookresearch/dino:main', 'dino_vitb16')

    if args.warmup_model_dir is not None:
        args.logger.info(f'Loading weights from {args.warmup_model_dir}')
        backbone.load_state_dict(torch.load(args.warmup_model_dir, map_location='cpu'))
    
    # NOTE: Hardcoded image size as we do not finetune the entire ViT model
    args.image_size = 224
    args.feat_dim = 768
    args.num_mlp_layers = 3
    args.mlp_out_dim = args.num_labeled_classes + args.num_unlabeled_classes
    if args.dataset_name == 'inat21_incremental':
        args.mlp_out_dim = 10000
    if args.init_estimate_num_cls is not None:
        args.mlp_out_dim = args.init_estimate_num_cls

    # ----------------------
    # HOW MUCH OF BASE MODEL TO FINETUNE
    # ----------------------
    for m in backbone.parameters():
        m.requires_grad = False

    # Only finetune layers from block 'args.grad_from_block' onwards
    for name, m in backbone.named_parameters():
        if 'block' in name:
            block_num = int(name.split('.')[1])
            if block_num >= args.grad_from_block:
                m.requires_grad = True

    train_transform, test_transform = get_transform(args.transform, image_size=args.image_size, args=args)
    train_transform = ContrastiveLearningViewGenerator(base_transform=train_transform, n_views=args.n_views)

    train_loaders, test_loaders_unlabelled, test_loaders_labelled = [], [], []

    wandb.init(project='IncreNat', entity=args.entity, name=args.exp_name, 
                tags=args.tags.split('-') if args.tags is not None else None, 
                resume=True if args.resume is not None else None)
    wandb.config.update(args)
                
    # get the correct names for args
    args.stage = 1
    _ = get_datasets(args.dataset_name, None, None, args)

    projector = DINOHead(in_dim=args.feat_dim, out_dim=args.mlp_out_dim, nlayers=args.num_mlp_layers)
    # Reload the model to be icarl
    model = iCarlNet(backbone, projector, args).to(device)
    args.logger.info('model build')
    
    args.classes_in_labelled_list = []
    args.classes_in_unlabelled_list = []
    for stage_idx in range(1, args.max_stage + 1):
        args.stage = stage_idx
        train_dataset, test_dataset, unlabelled_train_examples_test, _ = get_datasets(args.dataset_name,
                                                                                             train_transform,
                                                                                             test_transform,
                                                                                             args)
        # update train
        args.target_transform = train_dataset.labelled_dataset.target_transform
        args.classes_in_labelled_list.append(args.classes_in_labelled)
        args.classes_in_unlabelled_list.append(args.classes_in_unlabelled)

        if stage_idx > 1:
            model.construct_exemplar_sets(train_loaders[-1], args.exemplar_number, args)
            new_labelled_dataset = model.combine_dataset_with_exemplars(train_dataset.labelled_dataset, args.exemplar_number, args)
            new_labelled_dataset.uq_idxs = np.array(range(len(new_labelled_dataset)))
            
            train_dataset.labelled_dataset = new_labelled_dataset


        label_len = len(train_dataset.labelled_dataset)
        unlabelled_len = len(train_dataset.unlabelled_dataset)
        args.logger.info(f'label length: {label_len}, unlabel length: {unlabelled_len}')
        sample_weights = [1 if i < label_len else label_len / unlabelled_len for i in range(len(train_dataset))]
        sample_weights = torch.DoubleTensor(sample_weights)
        sampler = torch.utils.data.WeightedRandomSampler(sample_weights, num_samples=len(train_dataset))

        train_loader = DataLoader(train_dataset, num_workers=args.num_workers, batch_size=args.batch_size, shuffle=False,
                                  sampler=sampler, drop_last=True, pin_memory=True)
        test_loader_unlabelled = DataLoader(unlabelled_train_examples_test, num_workers=args.num_workers,
                                            batch_size=256, shuffle=False, pin_memory=False)
        test_loader_labelled = DataLoader(test_dataset, num_workers=args.num_workers,
                                          batch_size=256, shuffle=False, pin_memory=False)

        train_loaders.append(train_loader)
        test_loaders_unlabelled.append(test_loader_unlabelled)
        test_loaders_labelled.append(test_loader_labelled)

        train(model, train_loaders, test_loaders_labelled, test_loaders_unlabelled, args)

model/simgcd_icarl.py
DTennant/IGCD
import argparse

import math
import numpy as np
import torch
import torch.nn as nn
from torch.optim import SGD, lr_scheduler
from torch.utils.data import DataLoader
from tqdm import tqdm
import wandb
import inspect

from data.augmentations import get_transform
from data.get_datasets import get_datasets, get_class_splits

from util.general_utils import AverageMeter, init_experiment
from util.cluster_and_log_utils import log_accs_from_preds
from config import exp_root
from model.loss import info_nce_logits, SupConLoss, DistillLoss, ContrastiveLearningViewGenerator, get_params_groups


class DINOHead(nn.Module):
    def __init__(self, in_dim, out_dim, use_bn=False, norm_last_layer=True, 
                 nlayers=3, hidden_dim=2048, bottleneck_dim=256):
        super().__init__()
        nlayers = max(nlayers, 1)
        if nlayers == 1:
            self.mlp = nn.Linear(in_dim, bottleneck_dim)
        elif nlayers != 0:
            layers = [nn.Linear(in_dim, hidden_dim)]
            if use_bn:
                layers.append(nn.BatchNorm1d(hidden_dim))
            layers.append(nn.GELU())
            for _ in range(nlayers - 2):
                layers.append(nn.Linear(hidden_dim, hidden_dim))
                if use_bn:
                    layers.append(nn.BatchNorm1d(hidden_dim))
                layers.append(nn.GELU())
            layers.append(nn.Linear(hidden_dim, bottleneck_dim))
            self.mlp = nn.Sequential(*layers)
        self.apply(self._init_weights)
        self.last_layer = nn.utils.weight_norm(nn.Linear(in_dim, out_dim, bias=False))
        self.last_layer.weight_g.data.fill_(1)
        if norm_last_layer:
            self.last_layer.weight_g.requires_grad = False

    def _init_weights(self, m):
        if isinstance(m, nn.Linear):
            torch.nn.init.trunc_normal_(m.weight, std=.02)
            if isinstance(m, nn.Linear) and m.bias is not None:
                nn.init.constant_(m.bias, 0)

    def forward(self, x):
        ori_x = x
        x_proj = self.mlp(x)
        x = nn.functional.normalize(x, dim=-1, p=2)
        # x = x.detach()
        logits = self.last_layer(x)
        return x_proj, logits, ori_x


def train(student, train_loaders, test_loaders, unlabelled_train_loaders, args):
    params_groups = get_params_groups(student)
    optimizer = SGD(params_groups, lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay)
    fp16_scaler = None
    if args.fp16:
        fp16_scaler = torch.cuda.amp.GradScaler()

    exp_lr_scheduler = lr_scheduler.CosineAnnealingLR(
            optimizer,
            T_max=len(train_loaders[-1]) * (args.epochs - args.start_epoch),
            eta_min=args.lr * 1e-3,
        )

    if args.resume:
        load_dict = torch.load(args.resume, map_location='cpu')
        model_dict = load_dict['model']
        optimizer_dict = load_dict['optimizer']
        lr_scheduler_dict = load_dict['lr_scheduler']
        args.start_epoch = load_dict['epoch']
        
        student.load_state_dict(model_dict)
        optimizer.load_state_dict(optimizer_dict)
        exp_lr_scheduler.load_state_dict(lr_scheduler_dict)
        
        args.logger.info(f"Resumed from {args.resume}")

    cluster_criterion = DistillLoss(
                        args.warmup_teacher_temp_epochs,
                        args.epochs,
                        args.n_views,
                        args.warmup_teacher_temp,
                        args.teacher_temp,
                    )

    # inductive
    best_test_acc_ubl = 0
    # transductive
    best_train_acc_lab = 0
    best_train_acc_ubl = 0 
    best_train_acc_all = 0

    for epoch in range(args.start_epoch, args.epochs):
        loss_record = AverageMeter()

        student.train()
        for batch_idx, batch in enumerate(train_loaders[-1]):
            images, class_labels, uq_idxs, mask_lab = batch
            # mask_lab = mask_lab[:, 0]

            class_labels, mask_lab = class_labels.cuda(non_blocking=True), mask_lab.cuda(non_blocking=True).bool()
            images = torch.cat(images, dim=0).cuda(non_blocking=True)

            with torch.cuda.amp.autocast(fp16_scaler is not None):
                student_proj, student_out, _ = student(images)
                teacher_out = student_out.detach()

                # clustering, sup
                sup_logits = torch.cat([f[mask_lab] for f in (student_out / 0.1).chunk(2)], dim=0)
                sup_labels = torch.cat([class_labels[mask_lab] for _ in range(2)], dim=0)
                cls_loss = nn.CrossEntropyLoss()(sup_logits, sup_labels)

                # clustering, unsup
                cluster_loss = cluster_criterion(student_out, teacher_out, epoch)
                avg_probs = (student_out / 0.1).softmax(dim=1).mean(dim=0)
                me_max_loss = - torch.sum(torch.log(avg_probs**(-avg_probs))) + math.log(float(len(avg_probs)))
                cluster_loss += args.memax_weight * me_max_loss

                # represent learning, unsup
                contrastive_logits, contrastive_labels = info_nce_logits(features=student_proj)
                contrastive_loss = torch.nn.CrossEntropyLoss()(contrastive_logits, contrastive_labels)

                # representation learning, sup
                student_proj = torch.cat([f[mask_lab].unsqueeze(1) for f in student_proj.chunk(2)], dim=1)
                student_proj = torch.nn.functional.normalize(student_proj, dim=-1)
                sup_con_labels = class_labels[mask_lab]
                sup_con_loss = SupConLoss()(student_proj, labels=sup_con_labels)

                pstr = ''
                pstr += f'cls_loss: {cls_loss.item():.4f} '
                pstr += f'cluster_loss: {cluster_loss.item():.4f} '
                pstr += f'sup_con_loss: {sup_con_loss.item():.4f} '
                pstr += f'contrastive_loss: {contrastive_loss.item():.4f} '

                loss = 0
                loss += (1 - args.sup_weight) * cluster_loss + args.sup_weight * cls_loss
                loss += (1 - args.sup_weight) * contrastive_loss + args.sup_weight * sup_con_loss
                
            # Train acc
            loss_record.update(loss.item(), class_labels.size(0))
            optimizer.zero_grad()
            if fp16_scaler is None:
                loss.backward()
                optimizer.step()
            else:
                fp16_scaler.scale(loss).backward()
                fp16_scaler.step(optimizer)
                fp16_scaler.update()
            # Step schedule
            exp_lr_scheduler.step()

            if batch_idx % args.print_freq == 0:
                args.logger.info('Epoch: [{}][{}/{}]\t loss {:.5f}\t {}'
                            .format(epoch, batch_idx, len(train_loaders[-1]), loss.item(), pstr))
                
            wandb.log({
                'cls_loss': cls_loss.item(),
                'cluster_loss': cluster_loss.item(),
                'sup_con_loss': sup_con_loss.item(),
                'contrastive_loss': contrastive_loss.item(),
                'loss': loss.item(),
                'lr': optimizer.param_groups[0]['lr'],
            })

        args.logger.info('Train Epoch: {} Avg Loss: {:.4f} '.format(epoch, loss_record.avg))

        save_dict = {
            'model': student.state_dict(),
            'optimizer': optimizer.state_dict(),
            'lr_scheduler': exp_lr_scheduler.state_dict(),
            'epoch': epoch,
        }
        torch.save(save_dict, args.model_path[:-3] + f'_stage_{len(train_loaders)}.pth')
        args.logger.info("model saved to {}.".format(args.model_path[:-3] + f'_stage_{len(train_loaders)}.pth'))

        # args.logger.info('Testing on unlabelled examples in the training data...')
        # outputs = test_multi_test_loader(student, unlabelled_train_loaders, epoch=epoch, save_name='Train ACC Unlabelled', args=args)

        # args.logger.info('Testing on disjoint test set...')
        # outputs_test = test_multi_test_loader(student, test_loaders, epoch=epoch, save_name='Test ACC', args=args)

        for stage_idx, (unlabelled_train_loader, test_loader, classes_in_labelled, classes_in_unlabelled, ) in enumerate(zip(unlabelled_train_loaders, test_loaders, 
                                                                                                    args.classes_in_labelled_list, args.classes_in_unlabelled_list, )):
            # NOTE: here we need make sure these classes are in the target transform
            args.classes_in_labelled, args.classes_in_unlabelled = classes_in_labelled, classes_in_unlabelled
            stage_idx += 1

            args.logger.info('Testing on unlabelled examples in the training data...')
            output = test(student, unlabelled_train_loader, epoch=epoch, save_name=f'Train ACC Unlabelled stage {stage_idx}', stage_idx=stage_idx-1, args=args) 

            args.logger.info('Testing on disjoint test set...')
            output_test = test(student, test_loader, epoch=epoch, save_name=f'Test ACC stage {stage_idx}', stage_idx=stage_idx-1, args=args)

            all_acc, old_acc, new_acc = output['All'], output['Old'], output['New']
            all_acc_test, old_acc_test, new_acc_test = output_test['All'], output_test['Old'], output_test['New']
            
            wandb.log({f'stage_{stage_idx}_{len(train_loaders)}_' + k + '_train': v for k, v in output.items()})
            wandb.log({f'stage_{stage_idx}_{len(train_loaders)}_' + k + '_test': v for k, v in output_test.items()})

            args.logger.info('Train Accuracies stage {}: All {:.4f} | Old {:.4f} | New {:.4f}'.format(stage_idx, all_acc, old_acc, new_acc))
            args.logger.info('Test Accuracies stage {}: All {:.4f} | Old {:.4f} | New {:.4f}'.format(stage_idx, all_acc_test, old_acc_test, new_acc_test))

            if stage_idx == len(train_loaders):
                if new_acc_test > best_test_acc_ubl:

                    args.logger.info(f'Best ACC on old Classes on disjoint test set: {old_acc_test:.4f}...')
                    args.logger.info('Best Train Accuracies: All {:.4f} | Old {:.4f} | New {:.4f}'.format(all_acc, old_acc, new_acc))

                    torch.save(save_dict, args.model_path[:-3] + f'_{stage_idx}_{len(train_loaders)}_best.pt')
                    args.logger.info("model saved to {}.".format(args.model_path[:-3] + f'_{stage_idx}_best.pt'))

                    # inductive
                    best_output_test = output_test
                    best_test_acc_ubl = new_acc_test
                    # transductive            
                    best_output = output
                    best_train_acc_lab = old_acc
                    best_train_acc_ubl = new_acc
                    best_train_acc_all = all_acc

                args.logger.info(f'Exp Name: {args.exp_name} Exp id: {args.exp_id}')
                best_pstr = f'Metrics with best model on test set: '
                output_dict = best_output_test if args.eval_setting == 'inductive' else best_output
                for k, v in output_dict.items():
                    best_pstr += f'{k}: {v:.4f} '
                    wandb.log({f'stage_{stage_idx}_{len(train_loaders)}_best_{k}': v})
                args.logger.info(best_pstr)

def test(model, test_loader, epoch, save_name, stage_idx, args):
    model.eval()
    cur_stage = args.stage - 1

    preds, targets = [], []
    mask = np.array([])
    for batch_idx, (images, label, _) in enumerate(tqdm(test_loader)):
        images = images.cuda(non_blocking=True)
        with torch.no_grad():
            _, logits, _ = model(images)
            preds.append(logits.argmax(1).cpu().numpy())
            targets.append(label.cpu().numpy())
            mask = np.append(mask, np.array([True if x.item() in range(len(args.train_classes)) else False for x in label]))

    preds = np.concatenate(preds)
    targets = np.concatenate(targets)
    return log_accs_from_preds(y_true=targets, y_pred=preds, mask=mask,
                                                    T=epoch, eval_funcs=args.eval_funcs, save_name=save_name,
                                                    args=args)


if __name__ == "__main__":

    parser = argparse.ArgumentParser(description='cluster', formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument('--batch_size', default=128, type=int)
    parser.add_argument('--num_workers', default=8, type=int)
    parser.add_argument('--eval_funcs', nargs='+', help='Which eval functions to use', default=['v2', ])

    parser.add_argument('--warmup_model_dir', type=str, default=None)
    parser.add_argument('--resume', type=str, default=None)
    
    parser.add_argument('--n_overlap_cls', type=int, default=-1, help='how many classes overlaps, -1 means all labelled cls are in unlabelled')
    parser.add_argument('--dataset_name', type=str, default='scars', help='options: cifar10, cifar100, imagenet_100, cub, scars, fgvc_aricraft, herbarium_19')
    parser.add_argument('--prop_train_labels', type=float, default=0.5)
    parser.add_argument('--use_ssb_splits', action='store_true', default=True)
    parser.add_argument('--eval_setting', type=str, default='transductive', choices=['transductive', 'inductive'], help='options: inductive, transductive')

    parser.add_argument('--grad_from_block', type=int, default=11)
    parser.add_argument('--lr', type=float, default=0.1)
    parser.add_argument('--gamma', type=float, default=0.1)
    parser.add_argument('--momentum', type=float, default=0.9)
    parser.add_argument('--weight_decay', type=float, default=1e-4)
    parser.add_argument('--start_epoch', default=0, type=int)
    parser.add_argument('--epochs', default=200, type=int)
    parser.add_argument('--exp_root', type=str, default=exp_root)
    parser.add_argument('--transform', type=str, default='imagenet')
    parser.add_argument('--sup_weight', type=float, default=0.35)
    parser.add_argument('--n_views', default=2, type=int)
    
    parser.add_argument('--init_estimate_num_cls', type=int, help='the initial estimate of the number of classes, None if use GT', default=None)
    parser.add_argument('--memax_weight', type=float, default=2)
    parser.add_argument('--warmup_teacher_temp', default=0.07, type=float, help='Initial value for the teacher temperature.')
    parser.add_argument('--teacher_temp', default=0.04, type=float, help='Final value (after linear warmup)of the teacher temperature.')
    parser.add_argument('--warmup_teacher_temp_epochs', default=30, type=int, help='Number of warmup epochs for the teacher temperature.')

    parser.add_argument('--fp16', action='store_true', default=False)

    parser.add_argument('--use_small_set', action='store_true', default=False)
    parser.add_argument('--split_crit', type=str, )

    parser.add_argument('--return_path', action='store_true', default=False)
    
    parser.add_argument('--print_freq', default=10, type=int)
    parser.add_argument('--runner_name', default='simgcd', type=str, )
    parser.add_argument('--exp_name', default=None, type=str)
    parser.add_argument('--exp_id', default=None, type=str)
    parser.add_argument('--tags', default=None, type=str, help='- separated tags for the experiment; e.g. baseline-ssb')

    # ----------------------
    # INIT
    # ----------------------
    args = parser.parse_args()
    device = torch.device('cuda:0')
    args = get_class_splits(args)

    args.num_labeled_classes = len(args.train_classes)
    args.num_unlabeled_classes = len(args.unlabeled_classes)

    init_experiment(args, runner_name=[args.runner_name], exp_id=args.exp_id)
    args.logger.info(f'Using evaluation function {args.eval_funcs[0]} to print results')
    
    torch.backends.cudnn.benchmark = True

    if args.split_crit in ['year', 'loc_year']:
        args.max_stage = 4
    elif args.split_crit in ['location']:
        args.max_stage = 3
        

    # ----------------------
    # BASE MODEL
    # ----------------------
    args.interpolation = 3
    args.crop_pct = 0.875

    backbone = torch.hub.load('facebookresearch/dino:main', 'dino_vitb16')

    if args.warmup_model_dir is not None:
        args.logger.info(f'Loading weights from {args.warmup_model_dir}')
        backbone.load_state_dict(torch.load(args.warmup_model_dir, map_location='cpu'))
    
    # NOTE: Hardcoded image size as we do not finetune the entire ViT model
    args.image_size = 224
    args.feat_dim = 768
    args.num_mlp_layers = 3
    args.mlp_out_dim = args.num_labeled_classes + args.num_unlabeled_classes
    if args.dataset_name == 'inat21_incremental':
        args.mlp_out_dim = 10000
    if args.init_estimate_num_cls is not None:
        args.mlp_out_dim = args.init_estimate_num_cls

    # ----------------------
    # HOW MUCH OF BASE MODEL TO FINETUNE
    # ----------------------
    for m in backbone.parameters():
        m.requires_grad = False

    # Only finetune layers from block 'args.grad_from_block' onwards
    for name, m in backbone.named_parameters():
        if 'block' in name:
            block_num = int(name.split('.')[1])
            if block_num >= args.grad_from_block:
                m.requires_grad = True

    projector = DINOHead(in_dim=args.feat_dim, out_dim=args.mlp_out_dim, nlayers=args.num_mlp_layers)
    model = nn.Sequential(backbone, projector).to(device)
    args.logger.info('model build')
    
    wandb.init(project='IncreNat', name=args.exp_name, tags=args.tags.split('-') if args.tags is not None else None, resume=True if args.resume is not None else None)
    wandb.config.update(args)
    
    # --------------------
    # CONTRASTIVE TRANSFORM
    # --------------------
    train_transform, test_transform = get_transform(args.transform, image_size=args.image_size, args=args)
    train_transform = ContrastiveLearningViewGenerator(base_transform=train_transform, n_views=args.n_views)

    train_loaders, test_loaders_unlabelled, test_loaders_labelled = [], [], []
    
    # import ipdb; ipdb.set_trace()
    args.classes_in_labelled_list = []
    args.classes_in_unlabelled_list = []
    args.target_transforms = []
    for stage_idx in range(1, args.max_stage + 1):
        args.stage = stage_idx
        train_dataset, test_dataset, unlabelled_train_examples_test, _ = get_datasets(args.dataset_name,
                                                                                             train_transform,
                                                                                             test_transform,
                                                                                             args)
        # update train
        args.target_transform = train_dataset.labelled_dataset.target_transform
        args.target_transforms.append(args.target_transform)
        args.classes_in_labelled_list.append(args.classes_in_labelled)
        args.classes_in_unlabelled_list.append(args.classes_in_unlabelled)
        
        label_len = len(train_dataset.labelled_dataset)
        unlabelled_len = len(train_dataset.unlabelled_dataset)
        args.logger.info(f'label length: {label_len}, unlabel length: {unlabelled_len}')
        sample_weights = [1 if i < label_len else label_len / unlabelled_len for i in range(len(train_dataset))]
        sample_weights = torch.DoubleTensor(sample_weights)
        sampler = torch.utils.data.WeightedRandomSampler(sample_weights, num_samples=len(train_dataset))

        train_loader = DataLoader(train_dataset, num_workers=args.num_workers, batch_size=args.batch_size, shuffle=False,
                                  sampler=sampler, drop_last=True, pin_memory=True)
        test_loader_unlabelled = DataLoader(unlabelled_train_examples_test, num_workers=args.num_workers,
                                            batch_size=256, shuffle=False, pin_memory=False)
        test_loader_labelled = DataLoader(test_dataset, num_workers=args.num_workers,
                                          batch_size=256, shuffle=False, pin_memory=False)

        train_loaders.append(train_loader)
        test_loaders_unlabelled.append(test_loader_unlabelled)
        test_loaders_labelled.append(test_loader_labelled)

        train(model, train_loaders, test_loaders_labelled, test_loaders_unlabelled, args)

model/simgcd_incremental.py
DTennant/IGCD
import numpy as np
from scipy.optimize import linear_sum_assignment as linear_assignment

from sklearn import metrics
from sklearn.metrics.cluster import normalized_mutual_info_score as nmi_score
from sklearn.metrics import adjusted_rand_score as ari_score

def cluster_acc(y_true, y_pred, return_ind=False):
    """
    Calculate clustering accuracy. Require scikit-learn installed
    # Arguments
        y: true labels, numpy.array with shape `(n_samples,)`
        y_pred: predicted labels, numpy.array with shape `(n_samples,)`
    # Return
        accuracy, in [0,1]
    """
    # __import__("ipdb").set_trace()
    y_true = y_true.astype(int)
    assert y_pred.size == y_true.size
    D = max(y_pred.max(), y_true.max()) + 1
    w = np.zeros((D, D), dtype=int)
    for i in range(y_pred.size):
        w[y_pred[i], y_true[i]] += 1

    ind = linear_assignment(w.max() - w)
    ind = np.vstack(ind).T

    if return_ind:
        return sum([w[i, j] for i, j in ind]) * 1.0 / y_pred.size, ind, w
    else:
        return sum([w[i, j] for i, j in ind]) * 1.0 / y_pred.size


def purity_score(y_true, y_pred):
    # compute contingency matrix (also called confusion matrix)
    contingency_matrix = metrics.cluster.contingency_matrix(y_true, y_pred)
    # return purity
    return np.sum(np.amax(contingency_matrix, axis=0)) / np.sum(contingency_matrix)


def split_cluster_acc_v2(y_true, y_pred, mask, args=None):
    """
    Calculate clustering accuracy. Require scikit-learn installed
    First compute linear assignment on all data, then look at how good the accuracy is on subsets

    # Arguments
        mask: Which instances come from old classes (True) and which ones come from new classes (False)
        y: true labels, numpy.array with shape `(n_samples,)`
        y_pred: predicted labels, numpy.array with shape `(n_samples,)`
        args: arguments, not used

    # Return
        accuracy, in [0,1]
    """
    y_true = y_true.astype(int)

    old_classes_gt = set(y_true[mask])
    new_classes_gt = set(y_true[~mask])

    assert y_pred.size == y_true.size
    D = max(y_pred.max(), y_true.max()) + 1
    w = np.zeros((D, D), dtype=int)
    for i in range(y_pred.size):
        w[y_pred[i], y_true[i]] += 1

    ind = linear_assignment(w.max() - w)
    ind = np.vstack(ind).T

    ind_map = {j: i for i, j in ind}
    total_acc = sum([w[i, j] for i, j in ind]) * 1.0 / y_pred.size

    old_acc = 0
    total_old_instances = 0
    for i in old_classes_gt:
        old_acc += w[ind_map[i], i]
        total_old_instances += sum(w[:, i])
    old_acc /= total_old_instances

    new_acc = 0
    total_new_instances = 0
    for i in new_classes_gt:
        new_acc += w[ind_map[i], i]
        total_new_instances += sum(w[:, i])
    new_acc /= total_new_instances

    return {
        'All': total_acc, 
        'Old': old_acc, 
        'New': new_acc,
    }



def split_cluster_acc_v2_balanced(y_true, y_pred, mask, args=None):
    """
    Calculate clustering accuracy. Require scikit-learn installed
    First compute linear assignment on all data, then look at how good the accuracy is on subsets

    # Arguments
        mask: Which instances come from old classes (True) and which ones come from new classes (False)
        y: true labels, numpy.array with shape `(n_samples,)`
        y_pred: predicted labels, numpy.array with shape `(n_samples,)`
        args: arguments, not used

    # Return
        accuracy, in [0,1]
    """
    y_true = y_true.astype(int)

    old_classes_gt = set(y_true[mask])
    new_classes_gt = set(y_true[~mask])

    assert y_pred.size == y_true.size
    D = max(y_pred.max(), y_true.max()) + 1
    w = np.zeros((D, D), dtype=int)
    for i in range(y_pred.size):
        w[y_pred[i], y_true[i]] += 1

    ind = linear_assignment(w.max() - w)
    ind = np.vstack(ind).T

    ind_map = {j: i for i, j in ind}

    old_acc = np.zeros(len(old_classes_gt))
    total_old_instances = np.zeros(len(old_classes_gt))
    for idx, i in enumerate(old_classes_gt):
        old_acc[idx] += w[ind_map[i], i]
        total_old_instances[idx] += sum(w[:, i])

    new_acc = np.zeros(len(new_classes_gt))
    total_new_instances = np.zeros(len(new_classes_gt))
    for idx, i in enumerate(new_classes_gt):
        new_acc[idx] += w[ind_map[i], i]
        total_new_instances[idx] += sum(w[:, i])

    total_acc = np.concatenate([old_acc, new_acc]) / np.concatenate([total_old_instances, total_new_instances])
    old_acc /= total_old_instances
    new_acc /= total_new_instances
    total_acc[total_acc == np.inf] = 0
    old_acc[old_acc == np.inf] = 0
    new_acc[new_acc == np.inf] = 0
    total_acc, old_acc, new_acc = total_acc.mean(), old_acc.mean(), new_acc.mean()
    return {
        'All': total_acc, 
        'Old': old_acc, 
        'New': new_acc,
    }

def split_cluster_acc_v2_level(y_true, y_pred, mask, args=None):
    """
    Calculate clustering accuracy. Require scikit-learn installed
    First compute linear assignment on all data, then look at how good the accuracy is on subsets

    # Arguments
        mask: Which instances come from old classes (True) and which ones come from new classes (False)
        y: true labels, numpy.array with shape `(n_samples,)`
        y_pred: predicted labels, numpy.array with shape `(n_samples,)`
        args: arguments 

    # Return
        accuracy, in [0,1]
    """
    y_true = y_true.astype(int)

    old_classes_gt = set(y_true[mask])
    new_classes_gt = set(y_true[~mask])
    
    hard_classes_gt = set(args.target_transform(item) for item in args.hard_classes)
    medium_classes_gt = set(args.target_transform(item) for item in args.medium_classes)
    easy_classes_gt = set(args.target_transform(item) for item in args.easy_classes)

    assert y_pred.size == y_true.size
    D = max(y_pred.max(), y_true.max()) + 1
    w = np.zeros((D, D), dtype=int)
    for i in range(y_pred.size):
        w[y_pred[i], y_true[i]] += 1

    ind = linear_assignment(w.max() - w)
    ind = np.vstack(ind).T

    ind_map = {j: i for i, j in ind}
    
    old_acc = np.zeros(len(old_classes_gt))
    total_old_instances = np.zeros(len(old_classes_gt))
    for idx, i in enumerate(old_classes_gt):
        old_acc[idx] += w[ind_map[i], i]
        total_old_instances[idx] += sum(w[:, i])

    new_acc = np.zeros(len(new_classes_gt))
    total_new_instances = np.zeros(len(new_classes_gt))
    for idx, i in enumerate(new_classes_gt):
        new_acc[idx] += w[ind_map[i], i]
        total_new_instances[idx] += sum(w[:, i])

    total_acc = np.concatenate([old_acc, new_acc]) / np.concatenate([total_old_instances, total_new_instances])
    old_acc /= total_old_instances
    new_acc /= total_new_instances
    total_acc[total_acc == np.inf] = 0
    old_acc[old_acc == np.inf] = 0
    new_acc[new_acc == np.inf] = 0
    total_acc, old_acc, new_acc = total_acc.mean(), old_acc.mean(), new_acc.mean()

    # TODO: abstract this acc calculation
    hard_acc = np.zeros(len(hard_classes_gt))
    total_hard_instances = np.zeros(len(hard_classes_gt))
    for idx, i in enumerate(hard_classes_gt):
        hard_acc[idx] += w[ind_map[i], i]
        total_hard_instances[idx] += sum(w[:, i])

    medium_acc = np.zeros(len(medium_classes_gt))
    total_medium_instances = np.zeros(len(medium_classes_gt))
    for idx, i in enumerate(medium_classes_gt):
        medium_acc[idx] += w[ind_map[i], i]
        total_medium_instances[idx] += sum(w[:, i])
        
    easy_acc = np.zeros(len(easy_classes_gt))
    total_easy_instances = np.zeros(len(easy_classes_gt))
    for idx, i in enumerate(easy_classes_gt):
        easy_acc[idx] += w[ind_map[i], i]
        total_easy_instances[idx] += sum(w[:, i])

    hard_acc /= total_hard_instances
    medium_acc /= total_medium_instances
    easy_acc /= total_easy_instances
    hard_acc[hard_acc == np.inf] = 0
    medium_acc[medium_acc == np.inf] = 0
    easy_acc[easy_acc == np.inf] = 0
    hard_acc, medium_acc, easy_acc = hard_acc.mean(), medium_acc.mean(), easy_acc.mean()

    return {
        'All': total_acc, 
        'Old': old_acc, 
        'New': new_acc,
        'Hard': hard_acc,
        'Medium': medium_acc,
        'Easy': easy_acc,
    }

def split_cluster_acc_v2_overlap(y_true, y_pred, mask, args=None):
    """
    Calculate clustering accuracy. Require scikit-learn installed
    First compute linear assignment on all data, then look at how good the accuracy is on subsets

    # Arguments
        mask: Which instances come from old classes (True) and which ones come from new classes (False)
        y: true labels, numpy.array with shape `(n_samples,)`
        y_pred: predicted labels, numpy.array with shape `(n_samples,)`
        args: arguments 

    # Return
        accuracy, in [0,1]
    """
    y_true = y_true.astype(int)

    old_classes_gt = set(y_true[mask])
    new_classes_gt = set(y_true[~mask])
    
    overlap_classes_gt = set(args.target_transform(item) for item in args.overlap_cls)
    nonoverlap_classes_gt = set(args.target_transform(item) for item in set(args.train_classes) - set(args.overlap_cls))

    assert y_pred.size == y_true.size
    D = max(y_pred.max(), y_true.max()) + 1
    w = np.zeros((D, D), dtype=int)
    for i in range(y_pred.size):
        w[y_pred[i], y_true[i]] += 1

    ind = linear_assignment(w.max() - w)
    ind = np.vstack(ind).T

    ind_map = {j: i for i, j in ind}
    
    old_acc = np.zeros(len(old_classes_gt))
    total_old_instances = np.zeros(len(old_classes_gt))
    for idx, i in enumerate(old_classes_gt):
        old_acc[idx] += w[ind_map[i], i]
        total_old_instances[idx] += sum(w[:, i])

    new_acc = np.zeros(len(new_classes_gt))
    total_new_instances = np.zeros(len(new_classes_gt))
    for idx, i in enumerate(new_classes_gt):
        new_acc[idx] += w[ind_map[i], i]
        total_new_instances[idx] += sum(w[:, i])

    total_acc = np.concatenate([old_acc, new_acc]) / np.concatenate([total_old_instances, total_new_instances])
    old_acc /= total_old_instances
    new_acc /= total_new_instances
    total_acc[total_acc == np.inf] = 0
    old_acc[old_acc == np.inf] = 0
    new_acc[new_acc == np.inf] = 0
    total_acc, old_acc, new_acc = total_acc.mean(), old_acc.mean(), new_acc.mean()

    overlap_acc = np.zeros(len(overlap_classes_gt))
    total_overlap_instances = np.zeros(len(overlap_classes_gt))
    for idx, i in enumerate(overlap_classes_gt):
        overlap_acc[idx] += w[ind_map[i], i]
        total_overlap_instances[idx] += sum(w[:, i])

    nonoverlap_acc = np.zeros(len(nonoverlap_classes_gt))
    total_nonoverlap_instances = np.zeros(len(nonoverlap_classes_gt))
    for idx, i in enumerate(nonoverlap_classes_gt):
        nonoverlap_acc[idx] += w[ind_map[i], i]
        total_nonoverlap_instances[idx] += sum(w[:, i])
        
    overlap_acc /= total_overlap_instances
    try:
        nonoverlap_acc /= total_nonoverlap_instances
    except:
        nonoverlap_acc = np.array([0])
        args.logger.info('Using nonoverlap on training unlabelled set')
    overlap_acc[overlap_acc == np.inf] = 0
    nonoverlap_acc[nonoverlap_acc == np.inf] = 0
    overlap_acc, nonoverlap_acc = overlap_acc.mean(), nonoverlap_acc.mean()

    return {
        'All': total_acc, 
        'Old': old_acc, 
        'New': new_acc,
        'Overlap': overlap_acc,
        'Nonoverlap': nonoverlap_acc,
    }

def split_cluster_acc_v2_incremental(y_true, y_pred, mask, args=None):
    """
    Calculate clustering accuracy. Require scikit-learn installed
    First compute linear assignment on all data, then look at how good the accuracy is on subsets

    # Arguments
        mask: will be ignored
        y: true labels, numpy.array with shape `(n_samples,)`
        y_pred: predicted labels, numpy.array with shape `(n_samples,)`
        args: arguments 
            will use class_in_labelled and class_in_unlabelled

    # Return
        accuracy, in [0,1]
    """
    y_true = y_true.astype(int)
    
    # args.classes_in_labelled
    # args.classes_in_unlabelled
    
    # overlap_classes_gt = set(args.target_transform(item) for item in args.overlap_cls)
    # nonoverlap_classes_gt = set(args.target_transform(item) for item in set(args.train_classes) - set(args.overlap_cls))
    old_classes_gt = set(args.target_transform(item) for item in args.classes_in_labelled)
    new_classes_gt = set(args.target_transform(item) for item in args.classes_in_unlabelled)

    assert y_pred.size == y_true.size
    D = args.mlp_out_dim
    w = np.zeros((D, D), dtype=int)
    for i in range(y_pred.size):
        w[y_pred[i], y_true[i]] += 1

    ind = linear_assignment(w.max() - w)
    ind = np.vstack(ind).T

    ind_map = {j: i for i, j in ind}
    
    old_acc = np.zeros(D)
    total_old_instances = np.zeros(D)
    for idx, i in enumerate(old_classes_gt):
        old_acc[idx] += w[ind_map[i], i]
        total_old_instances[idx] += sum(w[:, i])

    new_acc = np.zeros(D)
    total_new_instances = np.zeros(D)
    for idx, i in enumerate(new_classes_gt):
        new_acc[idx] += w[ind_map[i], i]
        total_new_instances[idx] += sum(w[:, i])

    # total_acc = np.concatenate([old_acc, new_acc]) / np.concatenate([total_old_instances, total_new_instances])
    old_correct = np.array([correct / num_img for correct, num_img in zip(old_acc, total_old_instances) if num_img > 0])
    new_correct = np.array([correct / num_img for correct, num_img in zip(new_acc, total_new_instances) if num_img > 0])
    
    total_correct = np.array([correct / num_img for correct, num_img in zip(np.concatenate([old_acc, new_acc]), 
                                                         np.concatenate([total_old_instances + total_new_instances])) if num_img > 0])

    old_acc, new_acc = old_correct.mean(), new_correct.mean()
    total_acc = total_correct.mean()

    return {
        'All': total_acc, 
        'Old': old_acc, 
        'New': new_acc,
    }

def split_cluster_acc_v2_uni(y_true, y_pred, mask, args=None):
    """
    Calculate clustering accuracy. Require scikit-learn installed
    First compute linear assignment on all data, then look at how good the accuracy is on subsets

    # Arguments
        mask: Which instances come from old classes (True) and which ones come from new classes (False)
        y: true labels, numpy.array with shape `(n_samples,)`
        y_pred: predicted labels, numpy.array with shape `(n_samples,)`
        args: arguments 

    # Return
        accuracy, in [0,1]
    """
    y_true = y_true.astype(int)

    old_classes_gt = set(y_true[mask])
    new_classes_gt = set(y_true[~mask])
    
    hard_classes_gt = set(args.target_transform(item) for item in args.hard_classes)
    medium_classes_gt = set(args.target_transform(item) for item in args.medium_classes)
    easy_classes_gt = set(args.target_transform(item) for item in args.easy_classes)
    
    overlap_classes_gt = set(args.target_transform(item) for item in args.overlap_cls)
    nonoverlap_classes_gt = set(args.target_transform(item) for item in set(args.train_classes) - set(args.overlap_cls))

    assert y_pred.size == y_true.size
    D = max(y_pred.max(), y_true.max()) + 1
    w = np.zeros((D, D), dtype=int)
    for i in range(y_pred.size):
        w[y_pred[i], y_true[i]] += 1

    ind = linear_assignment(w.max() - w)
    ind = np.vstack(ind).T

    ind_map = {j: i for i, j in ind}
    
    old_acc = np.zeros(len(old_classes_gt))
    total_old_instances = np.zeros(len(old_classes_gt))
    for idx, i in enumerate(old_classes_gt):
        old_acc[idx] += w[ind_map[i], i]
        total_old_instances[idx] += sum(w[:, i])

    new_acc = np.zeros(len(new_classes_gt))
    total_new_instances = np.zeros(len(new_classes_gt))
    for idx, i in enumerate(new_classes_gt):
        new_acc[idx] += w[ind_map[i], i]
        total_new_instances[idx] += sum(w[:, i])

    total_acc = np.concatenate([old_acc, new_acc]) / np.concatenate([total_old_instances, total_new_instances])
    old_acc /= total_old_instances
    new_acc /= total_new_instances
    total_acc, old_acc, new_acc = total_acc.mean(), old_acc.mean(), new_acc.mean()

    # TODO: abstract this acc calculation
    hard_acc = np.zeros(len(hard_classes_gt))
    total_hard_instances = np.zeros(len(hard_classes_gt))
    for idx, i in enumerate(hard_classes_gt):
        hard_acc[idx] += w[ind_map[i], i]
        total_hard_instances[idx] += sum(w[:, i])

    medium_acc = np.zeros(len(medium_classes_gt))
    total_medium_instances = np.zeros(len(medium_classes_gt))
    for idx, i in enumerate(medium_classes_gt):
        medium_acc[idx] += w[ind_map[i], i]
        total_medium_instances[idx] += sum(w[:, i])
        
    easy_acc = np.zeros(len(easy_classes_gt))
    total_easy_instances = np.zeros(len(easy_classes_gt))
    for idx, i in enumerate(easy_classes_gt):
        easy_acc[idx] += w[ind_map[i], i]
        total_easy_instances[idx] += sum(w[:, i])

    hard_acc /= total_hard_instances
    medium_acc /= total_medium_instances
    easy_acc /= total_easy_instances
    hard_acc, medium_acc, easy_acc = hard_acc.mean(), medium_acc.mean(), easy_acc.mean()


    overlap_acc = np.zeros(len(overlap_classes_gt))
    total_overlap_instances = np.zeros(len(overlap_classes_gt))
    for idx, i in enumerate(overlap_classes_gt):
        overlap_acc[idx] += w[ind_map[i], i]
        total_overlap_instances[idx] += sum(w[:, i])

    nonoverlap_acc = np.zeros(len(nonoverlap_classes_gt))
    total_nonoverlap_instances = np.zeros(len(nonoverlap_classes_gt))
    for idx, i in enumerate(nonoverlap_classes_gt):
        nonoverlap_acc[idx] += w[ind_map[i], i]
        total_nonoverlap_instances[idx] += sum(w[:, i])
        
    overlap_acc /= total_overlap_instances
    try:
        nonoverlap_acc /= total_nonoverlap_instances
    except:
        nonoverlap_acc = np.array([0])
        args.logger.info('Using nonoverlap on training unlabelled set')
    overlap_acc, nonoverlap_acc = overlap_acc.mean(), nonoverlap_acc.mean()

    return {
        'All': total_acc, 
        'Old': old_acc, 
        'New': new_acc,
        'Hard': hard_acc,
        'Medium': medium_acc,
        'Easy': easy_acc,
        'Overlap': overlap_acc,
        'Nonoverlap': nonoverlap_acc,
    }

EVAL_FUNCS = {
    'v2': split_cluster_acc_v2,
    'v2b': split_cluster_acc_v2_balanced,
    'v2l': split_cluster_acc_v2_level,
    'v2o': split_cluster_acc_v2_overlap,
    'v2u': split_cluster_acc_v2_uni,
    'v2i': split_cluster_acc_v2_incremental,
}

def log_accs_from_preds(y_true, y_pred, mask, eval_funcs, save_name, T=None,
                        print_output=True, args=None):

    """
    Given a list of evaluation functions to use (e.g ['v1', 'v2']) evaluate and log ACC results

    :param y_true: GT labels
    :param y_pred: Predicted indices
    :param mask: Which instances belong to Old and New classes
    :param T: Epoch
    :param eval_funcs: Which evaluation functions to use
    :param save_name: What are we evaluating ACC on
    :param writer: Tensorboard logger
    :return:
    """

    mask = mask.astype(bool)
    y_true = y_true.astype(int)
    y_pred = y_pred.astype(int)

    for i, f_name in enumerate(eval_funcs):

        acc_f = EVAL_FUNCS[f_name]
        output_dict = acc_f(y_true, y_pred, mask, args)
        log_name = f'{save_name}_{f_name}'

        if i == 0:
            # to_return = (output_dict['All'], output_dict['Old'], output_dict['New'])
            to_return = output_dict

        if print_output:
            print_str = f'Epoch {T}, {log_name}: ' # All {all_acc:.4f} | Old {old_acc:.4f} | New {new_acc:.4f}'
            for item in output_dict:
                print_str += f'{item} {output_dict[item]:.4f} | '

            try:
                args.logger.info(print_str)
            except:
                print(print_str)

    return to_return
util/cluster_and_log_utils.py
DTennant/IGCD
import os
import torch
import inspect

from datetime import datetime
from loguru import logger

class AverageMeter(object):
    """Computes and stores the average and current value"""
    def __init__(self):
        self.reset()

    def reset(self):

        self.val = 0
        self.avg = 0
        self.sum = 0
        self.count = 0

    def update(self, val, n=1):

        self.val = val
        self.sum += val * n
        self.count += n
        self.avg = self.sum / self.count


def init_experiment(args, runner_name=None, exp_id=None):

    args.cuda = torch.cuda.is_available()

    # Get filepath of calling script
    if runner_name is None:
        runner_name = os.path.dirname(os.path.abspath(inspect.getfile(inspect.currentframe()))).split(".")[-2:]

    root_dir = os.path.join(args.exp_root, *runner_name)

    if not os.path.exists(root_dir):
        os.makedirs(root_dir)

    # Either generate a unique experiment ID, or use one which is passed
    if exp_id is None:

        if args.exp_name is None:
            raise ValueError("Need to specify the experiment name")
        # Unique identifier for experiment
        now = '{}_{:02d}.{:02d}.{}_'.format(args.exp_name, datetime.now().day, datetime.now().month, datetime.now().year) + \
              datetime.now().strftime("%S.%f")[:-3] 

        log_dir = os.path.join(root_dir, 'log', now)
        while os.path.exists(log_dir):
            now = '{:02d}.{:02d}.{}_'.format(datetime.now().day, datetime.now().month, datetime.now().year) + \
                  datetime.now().strftime("%S.%f")[:-3]

            log_dir = os.path.join(root_dir, 'log', now)

    else:

        log_dir = os.path.join(root_dir, 'log', f'{args.exp_name}_{exp_id}')

    if not os.path.exists(log_dir):
        os.makedirs(log_dir)
        
        
    logger.add(os.path.join(log_dir, 'log.txt'))
    args.logger = logger
    args.log_dir = log_dir

    # Instantiate directory to save models to
    model_root_dir = os.path.join(args.log_dir, 'checkpoints')
    if not os.path.exists(model_root_dir):
        os.mkdir(model_root_dir)

    args.model_dir = model_root_dir
    args.model_path = os.path.join(args.model_dir, 'model.pt')

    print(f'Experiment saved to: {args.log_dir}')

    hparam_dict = {}

    for k, v in vars(args).items():
        if isinstance(v, (int, float, str, bool, torch.Tensor)):
            hparam_dict[k] = v

    print(runner_name)
    print(args)

    return args

util/general_utils.py
DTennant/IGCD