experiments.py

import torch
import torch.nn as nn
from torch.nn.utils import weight_norm

from iCIFAR100 import iCIFAR100
import torchvision
import torch.optim as optim
import numpy as np
import copy
from torch.utils.data import DataLoader, Subset
from ResNet import resnet32
import torchvision.transforms as transforms

# TODO: da spostare
num_epochs = 70
batch_size = 128
DEVICE = 'cuda'
CIFAR100_MEAN = (0.5071, 0.4867, 0.4408)
CIFAR100_STD = (0.2675, 0.2565, 0.2761)
ROOT_FOLDER = "./data"


class ImbalancedDatasetSampler(torch.utils.data.sampler.Sampler):
    """Samples elements randomly from a given list of indices for imbalanced dataset
    Arguments:
        indices (list, optional): a list of indices
        num_samples (int, optional): number of samples to draw
        callback_get_label func: a callback-like function which takes two arguments - dataset and index
    """

    def __init__(self, dataset, n_classes, task_size, classes, indices=None, num_samples=None, callback_get_label=None):

        # if indices is not provided,
        # all elements in the dataset will be considered
        self.indices = list(range(len(dataset))) if indices is None else indices

        # define custom callback
        self.callback_get_label = callback_get_label

        # if num_samples is not provided,
        # draw `len(indices)` samples in each iteration
        self.n_classes = n_classes
        self.classes   = classes.tolist()
        self.task_size = task_size

        self.num_samples = len(self.indices) if num_samples is None else num_samples

        # distribution of classes in the dataset
        label_to_count = {}
        for idx in self.indices:
            label = self._get_label(dataset, idx)
            if label in label_to_count:
                label_to_count[label] += 1
            else:
                label_to_count[label] = 1

        weights = list()
        for idx in self.indices:
            label_idx = self._get_label(dataset, idx)
            if label_idx in self.classes[:self.n_classes - self.task_size]:
                weights.append(35) #7000/200
            elif label_idx in self.classes[self.n_classes - self.task_size:self.n_classes]:
                weights.append(14) #7000/500

        self.weights = torch.DoubleTensor(weights)

    def _get_label(self, dataset, idx):
        if isinstance(dataset, torchvision.datasets.CIFAR100):
            return dataset.targets[idx]
        elif isinstance(dataset, torch.utils.data.Subset):
            return dataset.dataset.targets[idx]
        elif self.callback_get_label:
            return self.callback_get_label(dataset, idx)
        else:
            raise NotImplementedError

    def __iter__(self):
        return (self.indices[i] for i in torch.multinomial(self.weights, self.num_samples, replacement=True))

    def __len__(self):
        return self.num_samples


class Network(nn.Module):
    def __init__(self, classifier, features_extractor):
        super(Network, self).__init__()
        self.classifier = classifier
        self.features_extractor = features_extractor
        self.fc = weight_norm(nn.Linear(features_extractor.fc.in_features, 100, bias=False), name="weight")
        torch.nn.init.kaiming_normal_(self.fc.weight, nonlinearity="relu")

    def forward(self, input):
        x = self.features_extractor(input)
        x = self.fc(x)
        return x


class iCarl:
    def __init__(self, n_classes, memory, classifier="NCM", examplar=True, random_exemplar=False):
        super(iCarl, self).__init__()

        """ Salvo l'istanza del modello attuale """
        self.model = Network(classifier, resnet32())

        """ Salvo il vecchio modello """
        self.old_model = None

        """ Numero di classi imparate ad un certo step """
        self.n_classes = n_classes

        """ Memoria assegnata (K) per salvare gli exemplar. """
        self.memory = memory

        """ Numero di classi che processo alla volta """
        self.task_size = 10

        """ Classificatore che uso. [NCM, KNN, SVM, MLP] """
        self.classifier = classifier

        """ True/False se fare uso o no di exemplar. """
        self.exemplar_usage = examplar

        """ lista delle medie per classe. """
        self.class_mean_set = []

        """ Scelgo gli exemplar casualmente o usando il criterio di icarl. """
        self.random_exemplar = random_exemplar

        """ Set degli exemplar, salvo solo gli indici """
        self.exemplar_sets = []

        """ Trasformazione per il training set"""
        self.train_transforms = transforms.Compose([transforms.RandomCrop((32, 32), padding=4),
                                                    transforms.RandomHorizontalFlip(p=0.5),
                                                    transforms.ToTensor(),
                                                    transforms.Normalize(CIFAR100_MEAN, CIFAR100_STD)])

        """ Trasformazioni per il test set """
        self.test_transforms = transforms.Compose([transforms.ToTensor(),
                                                   transforms.Normalize(CIFAR100_MEAN, CIFAR100_STD)])

        """ Trasformazioni per il classify """
        self.classify_transforms = transforms.Compose([transforms.RandomHorizontalFlip(p=1.),
                                                       transforms.ToTensor(),
                                                       transforms.Normalize(CIFAR100_MEAN, CIFAR100_STD)])
        """ Trasformazioni per il classify """
        self.exemplar_transforms_list = [transforms.RandomCrop((32, 32), padding=4),
                                         transforms.RandomHorizontalFlip(),
                                         transforms.RandomRotation(90),
                                         transforms.RandomVerticalFlip(),
                                         transforms.RandomGrayscale(),
                                         transforms.ColorJitter()]

        self.exemplar_transforms = transforms.Compose([transforms.RandomChoice(self.exemplar_transforms_list),
                                                       transforms.ToTensor(),
                                                       transforms.Normalize(CIFAR100_MEAN, CIFAR100_STD)])
        """ Dataset di train e test """
        self.train_dataset = iCIFAR100(ROOT_FOLDER,
                                       train=True,
                                       transform=self.train_transforms,
                                       t1=self.test_transforms,
                                       t2=self.classify_transforms,
                                       download=True)
        self.test_dataset = iCIFAR100(ROOT_FOLDER,
                                      train=False,
                                      transform=self.test_transforms,
                                      download=True)
        self.exemplar_dataset = iCIFAR100(ROOT_FOLDER,
                                          train=True,
                                          transform=self.exemplar_transforms,
                                          download=True)

        """ Dataloader """
        self.train_loader = None
        self.test_loader = None
        self.exemplar_loader = None

        """ Loss functions """
        self.BCE = nn.BCEWithLogitsLoss()

        """ Scelgo l'ordine delle 100 classi. """
        np.random.seed(1993)  # Fix the random seed
        self.classes = np.arange(100)
        np.random.shuffle(self.classes)

    def beforeTrain(self):
        """
        Procedure da eseguire prima del train come ad esempio incrementare il layer FC.
        :return: Void
        """
        self._update_dataloaders()
        self.model.to(DEVICE)

    def source_train_dataloader_iter(self):
        while True:
            for images, targets in self.exemplar_loader:
                yield images, targets

    def train(self):
        # Definisco l'optimizer e lo scheduler
        optimizer = optim.SGD(self.model.parameters(), lr=2., momentum=0.9, weight_decay=1e-5)
        scheduler = optim.lr_scheduler.MultiStepLR(optimizer, milestones=[48, 62], gamma=0.2)

        for epoch in range(num_epochs):
            running_loss = 0.0
            self.model.train()
            for images, labels in self.train_loader:
                images, labels = images.to(DEVICE), labels.to(DEVICE)
                optimizer.zero_grad()
                loss_value = self._compute_loss(images, labels)
                loss_value.backward()
                optimizer.step()
                running_loss += loss_value.item() * images.size(0)

            epoch_loss = running_loss / len(self.train_loader.dataset)
            accuracy = self._test()
            self.model.train()
            print('epoch:%d/%d,loss:%f,accuracy (FC layer):%.3f,LR=%s' % (epoch + 1, num_epochs, epoch_loss, accuracy, scheduler.get_last_lr()))
            scheduler.step()
        self.model.eval()
    """
    def train(self):
        # Definisco l'optimizer e lo scheduler
        optimizer = optim.SGD(self.model.parameters(), lr=2., momentum=0.9, weight_decay=1e-5)
        scheduler = optim.lr_scheduler.MultiStepLR(optimizer, milestones=[54, 70, 85], gamma=0.2) #48, 62

        for epoch in range(num_epochs):
            running_loss = 0.0
            self.model.train()
            if self.exemplar_loader is not None:
                gen1 = self.source_train_dataloader_iter()
            for images, labels in self.train_loader:
                if self.exemplar_loader is not None:
                    exemplar_img, exemplar_label = next(gen1)
                    exemplar_img, exemplar_label = exemplar_img.to(DEVICE), exemplar_label.to(DEVICE)
                images, labels = images.to(DEVICE), labels.to(DEVICE)
                optimizer.zero_grad()
                if self.exemplar_loader is not None:
                    loss_value = self._compute_loss(torch.cat((exemplar_img, images)), torch.cat((exemplar_label, labels)))
                else:
                    loss_value = self._compute_loss(images, labels)
                loss_value.backward()
                optimizer.step()
                running_loss += loss_value.item() * images.size(0)

            epoch_loss = running_loss / len(self.train_loader.dataset)
            accuracy = self._test()
            self.model.train()
            print('epoch:%d/%d,loss:%f,accuracy (FC layer):%.3f,LR=%s' % (epoch + 1, num_epochs, epoch_loss, accuracy, scheduler.get_last_lr()))
            scheduler.step()
        self.model.eval()
    """
    def afterTrain(self):
        """
        Eseguo le operazioni sugli exemplar ed eseguo il test finale.
        :return:
        """
        self.model.eval()
        m = self.memory // self.n_classes
        if self.exemplar_usage:
            self._compute_exemplar_class_mean()
            self._reduce_exemplar_sets(m)
            for i in self.classes[self.n_classes - self.task_size: self.n_classes]:
                print('construct class %s examplar:' % i, end='')
                images, indexes, _ = self.train_dataset.get_images_by_class(i)
                self._construct_exemplar_set(images, indexes, m)

        # self.model.train()
        accuracy = self._test(True)

        self.model.eval()
        self.old_model = Network(self.classifier, resnet32())
        self.old_model.load_state_dict(self.model.state_dict())
        self.old_model = self.old_model.to(DEVICE)
        # self.old_model = copy.deepcopy(self.model).to(DEVICE)
        self.old_model.eval()

        self.n_classes += self.task_size
        print(self.classifier + " accuracy：" + str(accuracy))

    def _test(self, final_step=False):

        self.model.eval()  # Set Network to evaluation mode

        running_corrects = 0

        for images, labels in self.test_loader:
            images = images.to(DEVICE)
            labels = labels.to(DEVICE)
            with torch.no_grad():
                if self.classifier == "NCM" and final_step:
                    preds = self.classify(images).to(DEVICE)
                else:
                    # Forward Pass
                    outputs = self.model(images)

                    # Get predictions
                    _, preds = torch.max(outputs.data, 1)

            # Update Corrects
            running_corrects += torch.sum(preds == labels.data).data.item()

        # Calculate Accuracy
        return 100 * running_corrects / float(len(self.test_loader.dataset))

    def classify(self, images):
        # batch_sizex3x32x32
        result = []
        self.model.eval()
        with torch.no_grad():
            phi_X = nn.functional.normalize(self.model.features_extractor(images))

        ex_means = torch.stack(self.class_mean_set)
        # 10x64 (di ogni classe mi salvo la media di ogni features)
        for x in phi_X:
            # x: 64. media delle features di quella immagine
            distances_from_class = (ex_means - x).norm(dim=1)  # giusto
            y = distances_from_class.argmin()
            result.append(self.classes[y])
        return torch.tensor(result)

    def _compute_exemplar_class_mean(self):
        self.class_mean_set = []
        self.model.eval()
        if self.n_classes > 10:
            for label, P_y in enumerate(self.exemplar_sets):
                images_1, images_2 = self.train_dataset.get_images_by_indexes(P_y)
                class_mean_1, _ = self.compute_class_mean(images_1)
                class_mean_2, _ = self.compute_class_mean(images_2)
                class_mean = (class_mean_1 + class_mean_2) / 2
                class_mean = class_mean.data / class_mean.norm()
                self.class_mean_set.append(class_mean)
        for i in self.classes[self.n_classes - self.task_size: self.n_classes]:
            images_1, _, images_2 = self.train_dataset.get_images_by_class(i)
            class_mean_1, _ = self.compute_class_mean(images_1)
            class_mean_2, _ = self.compute_class_mean(images_2)
            class_mean = (class_mean_1 + class_mean_2) / 2
            class_mean = class_mean.data / class_mean.norm()
            self.class_mean_set.append(class_mean)

    def _construct_exemplar_set(self, images, ind, m):
        """
        Costruisco il set degli exemplar basato sugli indici.
        :param images: tutte le immagini di quella classe
        :param m: numero di immagini da salvare
        :return:
        """
        self.model.eval()
        if self.random_exemplar:
            Py = list()
            indexes = np.arange(len(images))
            np.random.shuffle(indexes)
            for i in range(m):
                Py.append(ind[indexes[i]])
        else:
            images = torch.stack(images).to(DEVICE)  # 500x3x32x32
            with torch.no_grad():
                phi_X = torch.nn.functional.normalize(self.model.features_extractor(images)).cpu()  # 500x64

            mu_y = phi_X.mean(dim=0)  # vettore di 64 colonne
            mu_y.data = mu_y.data / mu_y.data.norm()

            Py = []
            # Accumulates sum of exemplars
            sum_taken_exemplars = torch.zeros(1, 64)

            indexes = list()
            for k in range(1, int(m + 1)):
                # Using broadcast: expanding mu_y and sum_taken_exemplars to phi_X shape
                asd = nn.functional.normalize((1 / k) * (phi_X + sum_taken_exemplars))
                mean_distances = (mu_y - asd).norm(dim=1)  # senza norma 500x1
                # min_index = mean_distances.argmin(dim=0).item()
                used = -1
                a, indici = torch.sort(mean_distances)
                for item in a:
                    mins = (mean_distances == item).nonzero()
                    for j in mins:
                        if j not in indexes:
                            indexes.append(j)
                            Py.append(ind[j])
                            used = j
                            sum_taken_exemplars += phi_X[j]
                            break
                    if used != -1:
                        break

        print(len(Py))
        # Py = torch.stack(Py) Lo tolgo visto che salvo gli indici
        self.exemplar_sets.append(Py)  # for dictionary version: self.exemplar_sets[y] = Py

    def compute_class_mean(self, images):
        """
        Passo tutte le immagini di una determinata classe e faccio la media.
        :param special_transform:
        :param images: tutte le immagini della classe x
        :return: media della classe e features extractor.
        """
        self.model.eval()
        images = torch.stack(images).to(DEVICE)  # 500x3x32x32  #stack vs cat. Il primo le attacca in una nuova dim. 3x4 diventa 1x3x4.
        # cat invece le fa diventare 6x4
        with torch.no_grad():
            phi_X = torch.nn.functional.normalize(self.model.features_extractor(images))

        # phi_X.shape = 500x64
        mean = phi_X.mean(dim=0)  # array 64. è la media di tutte le colonne
        mean.data = mean.data / mean.data.norm()
        return mean, phi_X

    def _reduce_exemplar_sets(self, images_per_class):
        for index in range(len(self.exemplar_sets)):
            self.exemplar_sets[index] = self.exemplar_sets[index][:images_per_class]
            print('Reduce size of class %d to %s examplar' % (self.classes[index], str(len(self.exemplar_sets[index]))))

    def _update_dataloaders(self):
        """
        Aggiorno i dataloader con le nuove immagini/labels delle nuove classi.
        :rtype: object
        """
        train_indexes = []
        if self.exemplar_usage:
            for i in self.exemplar_sets:
                train_indexes.extend(i)
        train_indexes.extend(self.train_dataset.get_indexes_by_classes(self.classes[self.n_classes - self.task_size: self.n_classes]))

        self.train_loader = DataLoader(dataset=self.train_dataset,
                                       shuffle=False,
                                       sampler=ImbalancedDatasetSampler(dataset=self.train_dataset,
                                                                        indices=train_indexes,
                                                                        n_classes=self.n_classes,
                                                                        task_size=self.task_size,
                                                                        classes=self.classes),
                                       num_workers=4,
                                       batch_size=128)
        """
        self.train_loader = DataLoader(Subset(self.train_dataset, train_indexes),
                                       shuffle=True,
                                       num_workers=4,
                                       batch_size=128)

        
        if len(exemplar_indexes) != 0:
            self.exemplar_loader = DataLoader(Subset(self.exemplar_dataset, exemplar_indexes),
                                              shuffle=True,
                                              num_workers=4,
                                              batch_size=128)
        """
        print(len(self.train_loader.dataset))
        test_indexes = self.test_dataset.get_indexes_by_classes(self.classes[:self.n_classes])
        self.test_loader = DataLoader(dataset=Subset(self.test_dataset, test_indexes),
                                      shuffle=False,
                                      num_workers=4,
                                      batch_size=128)
        print(len(self.test_loader.dataset))

    def _compute_loss(self, images, target):
        """
        Calcolo la loss usando la BCEWithLogits singola (senza usarne 2 separate)
        :param images: 128 immagini da processare
        :param target: 128 true labels
        :return: la loss
        """
        self.model.train()
        output = self.model(images)
        target = self.to_onehot(target, 100)
        output, target = output.to(DEVICE), target.to(DEVICE)
        if self.old_model is None:
            return self.BCE(output, target)
        else:
            with torch.no_grad():
                old_target = torch.sigmoid(self.old_model(images))

            n_c = self.classes[:self.n_classes - self.task_size]
            target[:, n_c] = old_target[:, n_c]
            return self.BCE(output, target)

    def _bce_bce_loss(self, images, target):
        self.model.train()
        output = self.model(images)
        target = self.to_onehot(target, 100)
        output, target = output.to(DEVICE), target.to(DEVICE)
        alfa = 0.6
        beta = 0.4
        loss = alfa * self.BCE(output, target)
        if self.old_model is not None:
            with torch.no_grad():
                old_target = torch.sigmoid(self.old_model(images))

            n_c = self.classes[:self.n_classes - self.task_size]
            loss += beta * self.BCE(output[:, n_c], old_target[:, n_c])

        return loss
    @staticmethod
    def to_onehot(targets, n_classes):
        return torch.eye(n_classes)[targets]