models/coil.py

import logging
import numpy as np
from tqdm import tqdm
import torch
from torch import optim
from torch.nn import functional as F
from torch.utils.data import DataLoader
from models.base import BaseLearner
from utils.inc_net import (
 IncrementalNet,
 CosineIncrementalNet,
 SimpleCosineIncrementalNet,
)
from utils.toolkit import target2onehot, tensor2numpy
import ot
from torch import nn
import copy

EPSILON = 1e-8

epochs = 100
lrate = 0.1
milestones = [40, 80]
lrate_decay = 0.1
batch_size = 32
memory_size = 2000
T = 2


class COIL(BaseLearner):
 def __init__(self, args):
 super().__init__(args)
 self._network = SimpleCosineIncrementalNet(args, False)
 self.data_manager = None
 self.nextperiod_initialization = None
 self.sinkhorn_reg = args["sinkhorn"]
 self.calibration_term = args["calibration_term"]
 self.args = args

 def after_task(self):
 self.nextperiod_initialization = self.solving_ot()
 self._old_network = self._network.copy().freeze()
 self._known_classes = self._total_classes

 def solving_ot(self):
 with torch.no_grad():
 if self._total_classes == self.data_manager.get_total_classnum():
 print("training over, no more ot solving")
 return None
 each_time_class_num = self.data_manager.get_task_size(1)
 self._extract_class_means(
 self.data_manager, 0, self._total_classes + each_time_class_num
 )
 former_class_means = torch.tensor(
 self._ot_prototype_means[: self._total_classes]
 )
 next_period_class_means = torch.tensor(
 self._ot_prototype_means[
 self._total_classes : self._total_classes + each_time_class_num
 ]
 )
 Q_cost_matrix = torch.cdist(
 former_class_means, next_period_class_means, p=self.args["norm_term"]
 )
 # solving ot
 _mu1_vec = (
 torch.ones(len(former_class_means)) / len(former_class_means) * 1.0
 )
 _mu2_vec = (
 torch.ones(len(next_period_class_means)) / len(former_class_means) * 1.0
 )
 T = ot.sinkhorn(_mu1_vec, _mu2_vec, Q_cost_matrix, self.sinkhorn_reg)
 T = torch.tensor(T).float().cuda()
 transformed_hat_W = torch.mm(
 T.T, F.normalize(self._network.fc.weight, p=2, dim=1)
 )
 oldnorm = torch.norm(self._network.fc.weight, p=2, dim=1)
 newnorm = torch.norm(
 transformed_hat_W * len(former_class_means), p=2, dim=1
 )
 meannew = torch.mean(newnorm)
 meanold = torch.mean(oldnorm)
 gamma = meanold / meannew
 self.calibration_term = gamma
 self._ot_new_branch = (
 transformed_hat_W * len(former_class_means) * self.calibration_term
 )
 return transformed_hat_W * len(former_class_means) * self.calibration_term

 def solving_ot_to_old(self):
 current_class_num = self.data_manager.get_task_size(self._cur_task)
 self._extract_class_means_with_memory(
 self.data_manager, self._known_classes, self._total_classes
 )
 former_class_means = torch.tensor(
 self._ot_prototype_means[: self._known_classes]
 )
 next_period_class_means = torch.tensor(
 self._ot_prototype_means[self._known_classes : self._total_classes]
 )
 Q_cost_matrix = (
 torch.cdist(
 next_period_class_means, former_class_means, p=self.args["norm_term"]
 )
 + EPSILON
 ) # in case of numerical err
 _mu1_vec = torch.ones(len(former_class_means)) / len(former_class_means) * 1.0
 _mu2_vec = (
 torch.ones(len(next_period_class_means)) / len(former_class_means) * 1.0
 )
 T = ot.sinkhorn(_mu2_vec, _mu1_vec, Q_cost_matrix, self.sinkhorn_reg)
 T = torch.tensor(T).float().cuda()
 transformed_hat_W = torch.mm(
 T.T,
 F.normalize(self._network.fc.weight[-current_class_num:, :], p=2, dim=1),
 )
 return transformed_hat_W * len(former_class_means) * self.calibration_term

 def incremental_train(self, data_manager):
 self._cur_task += 1
 self._total_classes = self._known_classes + data_manager.get_task_size(
 self._cur_task
 )

 self._network.update_fc(self._total_classes, self.nextperiod_initialization)
 self.data_manager = data_manager

 logging.info(
 "Learning on {}-{}".format(self._known_classes, self._total_classes)
 )
 self.lamda = self._known_classes / self._total_classes
 # Loader
 train_dataset = data_manager.get_dataset(
 np.arange(self._known_classes, self._total_classes),
 source="train",
 mode="train",
 appendent=self._get_memory(),
 )
 self.train_loader = DataLoader(
 train_dataset, batch_size=batch_size, shuffle=True, num_workers=4
 )
 test_dataset = data_manager.get_dataset(
 np.arange(0, self._total_classes), source="test", mode="test"
 )
 self.test_loader = DataLoader(
 test_dataset, batch_size=batch_size, shuffle=False, num_workers=4
 )

 self._train(self.train_loader, self.test_loader)
 
 if self.args['fixed_memory']:
 examplar_size = self.args["memory_per_class"]
 else:
 examplar_size = memory_size // self._total_classes
 self._reduce_exemplar(data_manager, examplar_size)
 self._construct_exemplar(data_manager, examplar_size)

 def _train(self, train_loader, test_loader):
 self._network.to(self._device)
 if self._old_network is not None:
 self._old_network.to(self._device)
 optimizer = optim.SGD(
 self._network.parameters(), lr=lrate, momentum=0.9, weight_decay=5e-4
 ) # 1e-5
 scheduler = optim.lr_scheduler.MultiStepLR(
 optimizer=optimizer, milestones=milestones, gamma=lrate_decay
 )
 self._update_representation(train_loader, test_loader, optimizer, scheduler)

 def _update_representation(self, train_loader, test_loader, optimizer, scheduler):
 prog_bar = tqdm(range(epochs))
 for _, epoch in enumerate(prog_bar):
 weight_ot_init = max(1.0 - (epoch / 2) ** 2, 0)
 weight_ot_co_tuning = (epoch / epochs) ** 2.0

 self._network.train()
 losses = 0.0
 correct, total = 0, 0

 for i, (_, inputs, targets) in enumerate(train_loader):
 inputs, targets = inputs.to(self._device), targets.to(self._device)
 output = self._network(inputs)
 logits = output["logits"]
 onehots = target2onehot(targets, self._total_classes)

 clf_loss = F.cross_entropy(logits, targets)
 if self._old_network is not None:

 old_logits = self._old_network(inputs)["logits"].detach()
 hat_pai_k = F.softmax(old_logits / T, dim=1)
 log_pai_k = F.log_softmax(
 logits[:, : self._known_classes] / T, dim=1
 )
 distill_loss = -torch.mean(torch.sum(hat_pai_k * log_pai_k, dim=1))

 if epoch < 1:
 features = F.normalize(output["features"], p=2, dim=1)
 current_logit_new = F.log_softmax(
 logits[:, self._known_classes :] / T, dim=1
 )
 new_logit_by_wnew_init_by_ot = F.linear(
 features, F.normalize(self._ot_new_branch, p=2, dim=1)
 )
 new_logit_by_wnew_init_by_ot = F.softmax(
 new_logit_by_wnew_init_by_ot / T, dim=1
 )
 new_branch_distill_loss = -torch.mean(
 torch.sum(
 current_logit_new * new_logit_by_wnew_init_by_ot, dim=1
 )
 )

 loss = (
 distill_loss * self.lamda
 + clf_loss * (1 - self.lamda)
 + 0.001 * (weight_ot_init * new_branch_distill_loss)
 )
 else:
 features = F.normalize(output["features"], p=2, dim=1)
 if i % 30 == 0:
 with torch.no_grad():
 self._ot_old_branch = self.solving_ot_to_old()
 old_logit_by_wold_init_by_ot = F.linear(
 features, F.normalize(self._ot_old_branch, p=2, dim=1)
 )
 old_logit_by_wold_init_by_ot = F.log_softmax(
 old_logit_by_wold_init_by_ot / T, dim=1
 )
 old_branch_distill_loss = -torch.mean(
 torch.sum(hat_pai_k * old_logit_by_wold_init_by_ot, dim=1)
 )
 loss = (
 distill_loss * self.lamda
 + clf_loss * (1 - self.lamda)
 + self.args["reg_term"]
 * (weight_ot_co_tuning * old_branch_distill_loss)
 )
 else:
 loss = clf_loss

 optimizer.zero_grad()
 loss.backward()
 optimizer.step()
 losses += loss.item()

 _, preds = torch.max(logits, dim=1)
 correct += preds.eq(targets.expand_as(preds)).cpu().sum()
 total += len(targets)

 scheduler.step()
 train_acc = np.around(tensor2numpy(correct) * 100 / total, decimals=2)
 test_acc = self._compute_accuracy(self._network, test_loader)
 info = "Task {}, Epoch {}/{} => Loss {:.3f}, Train_accy {:.2f}, Test_accy {:.2f}".format(
 self._cur_task,
 epoch + 1,
 epochs,
 losses / len(train_loader),
 train_acc,
 test_acc,
 )
 prog_bar.set_description(info)

 logging.info(info)

 def _extract_class_means(self, data_manager, low, high):
 self._ot_prototype_means = np.zeros(
 (data_manager.get_total_classnum(), self._network.feature_dim)
 )
 with torch.no_grad():
 for class_idx in range(low, high):
 data, targets, idx_dataset = data_manager.get_dataset(
 np.arange(class_idx, class_idx + 1),
 source="train",
 mode="test",
 ret_data=True,
 )
 idx_loader = DataLoader(
 idx_dataset, batch_size=batch_size, shuffle=False, num_workers=4
 )
 vectors, _ = self._extract_vectors(idx_loader)
 vectors = (vectors.T / (np.linalg.norm(vectors.T, axis=0) + EPSILON)).T
 class_mean = np.mean(vectors, axis=0)
 class_mean = class_mean / (np.linalg.norm(class_mean))
 self._ot_prototype_means[class_idx, :] = class_mean
 self._network.train()

 def _extract_class_means_with_memory(self, data_manager, low, high):

 self._ot_prototype_means = np.zeros(
 (data_manager.get_total_classnum(), self._network.feature_dim)
 )
 memoryx, memoryy = self._data_memory, self._targets_memory
 with torch.no_grad():
 for class_idx in range(0, low):
 idxes = np.where(
 np.logical_and(memoryy >= class_idx, memoryy < class_idx + 1)
 )[0]
 data, targets = memoryx[idxes], memoryy[idxes]
 # idx_dataset=TensorDataset(data,targets)
 # idx_loader = DataLoader(idx_dataset, batch_size=batch_size, shuffle=False, num_workers=4)
 _, _, idx_dataset = data_manager.get_dataset(
 [],
 source="train",
 appendent=(data, targets),
 mode="test",
 ret_data=True,
 )
 idx_loader = DataLoader(
 idx_dataset, batch_size=batch_size, shuffle=False, num_workers=4
 )
 vectors, _ = self._extract_vectors(idx_loader)
 vectors = (vectors.T / (np.linalg.norm(vectors.T, axis=0) + EPSILON)).T
 class_mean = np.mean(vectors, axis=0)
 class_mean = class_mean / np.linalg.norm(class_mean)
 self._ot_prototype_means[class_idx, :] = class_mean

 for class_idx in range(low, high):
 data, targets, idx_dataset = data_manager.get_dataset(
 np.arange(class_idx, class_idx + 1),
 source="train",
 mode="test",
 ret_data=True,
 )
 idx_loader = DataLoader(
 idx_dataset, batch_size=batch_size, shuffle=False, num_workers=4
 )
 vectors, _ = self._extract_vectors(idx_loader)
 vectors = (vectors.T / (np.linalg.norm(vectors.T, axis=0) + EPSILON)).T
 class_mean = np.mean(vectors, axis=0)
 class_mean = class_mean / np.linalg.norm(class_mean)
 self._ot_prototype_means[class_idx, :] = class_mean
 self._network.train()