models/beef_iso.py

import copy
import logging
import numpy as np
from tqdm import tqdm
import torch
from torch import nn
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 BEEFISONet
from utils.toolkit import count_parameters, target2onehot, tensor2numpy


EPSILON = 1e-8


class BEEFISO(BaseLearner):
 def __init__(self, args):
 super().__init__(args)
 self.args = args
 self._network = BEEFISONet(args, False)
 self._snet = None
 self.logits_alignment = args["logits_alignment"]
 self.val_loader = None
 self.reduce_batch_size = args["reduce_batch_size"]
 self.random = args.get("random",None)
 self.imbalance = args.get("imbalance",None)

 def after_task(self):
 self._network_module_ptr.update_fc_after()
 self._known_classes = self._total_classes
 if self.reduce_batch_size:
 if self._cur_task == 0:
 self.args["batch_size"] = self.args["batch_size"]
 else:
 self.args["batch_size"] = self.args["batch_size"] * (self._cur_task+1) // (self._cur_task+2) 
 logging.info("Exemplar size: {}".format(self.exemplar_size))

 def incremental_train(self, data_manager):
 self.data_manager = data_manager
 self._cur_task += 1
 if self._cur_task > 1 and self.args["is_compress"]:
 self._network = self._snet
 self._total_classes = self._known_classes + data_manager.get_task_size(
 self._cur_task
 )
 self._network.update_fc_before(self._total_classes)
 self._network_module_ptr = self._network
 logging.info(
 "Learning on {}-{}".format(self._known_classes, self._total_classes)
 )

 if self._cur_task > 0:
 for id in range(self._cur_task):
 for p in self._network.convnets[id].parameters():
 p.requires_grad = False
 for p in self._network.old_fc.parameters():
 p.requires_grad = False


 logging.info("All params: {}".format(count_parameters(self._network)))
 logging.info(
 "Trainable params: {}".format(count_parameters(self._network, True))
 )


 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=self.args["batch_size"],
 shuffle=True,
 num_workers=self.args["num_workers"],
 pin_memory=True,
 )
 test_dataset = data_manager.get_dataset(
 np.arange(0, self._total_classes), source="test", mode="test"
 )
 self.test_loader = DataLoader(
 test_dataset,
 batch_size=self.args["batch_size"],
 shuffle=False,
 num_workers=self.args["num_workers"],
 pin_memory=True,
 )
 if self._cur_task > 0:
 if self.random or self.imbalance:
 val_dset = data_manager.get_finetune_dataset(known_classes=self._known_classes, total_classes=self._total_classes,
 source="train", mode='train', appendent=self._get_memory(), type="ratio")
 else:
 _, val_dset = data_manager.get_dataset_with_split(np.arange(self._known_classes, self._total_classes),
 source='train', mode='train',
 appendent=self._get_memory(),
 val_samples_per_class=int(
 self.samples_old_class))
 self.val_loader = DataLoader(
 val_dset, batch_size=self.args["batch_size"], shuffle=True, num_workers=self.args["num_workers"], pin_memory=True)
 if len(self._multiple_gpus) > 1:
 self._network = nn.DataParallel(self._network, self._multiple_gpus)
 self._train(self.train_loader, self.test_loader,self.val_loader)
 if self.random or self.imbalance:
 self.build_rehearsal_memory_imbalance(data_manager,self.samples_per_class)
 else:
 self.build_rehearsal_memory(data_manager, self.samples_per_class)
 if len(self._multiple_gpus) > 1:
 self._network = self._network.module

 def train(self):
 self._network_module_ptr.train()
 self._network_module_ptr.convnets[-1].train()
 if self._cur_task >= 1:
 self._network_module_ptr.convnets[0].eval()

 def _train(self, train_loader, test_loader, val_loader=None):
 self._network.to(self._device)
 if hasattr(self._network, "module"):
 self._network_module_ptr = self._network.module
 if self._cur_task == 0:
 optimizer = optim.SGD(
 filter(lambda p: p.requires_grad, self._network.parameters()),
 momentum=0.9,
 lr=self.args["init_lr"],
 weight_decay=self.args["init_weight_decay"],
 )
 scheduler = optim.lr_scheduler.CosineAnnealingLR(
 optimizer=optimizer, T_max=self.args["init_epochs"]
 )
 self.epochs = self.args["init_epochs"]
 self._init_train(train_loader, test_loader, optimizer, scheduler)
 else:

 optimizer = optim.SGD(
 filter(lambda p: p.requires_grad, self._network.parameters()),
 lr=self.args["lr"],
 momentum=0.9,
 weight_decay=self.args["weight_decay"],
 )
 scheduler = optim.lr_scheduler.CosineAnnealingLR(
 optimizer=optimizer, T_max=self.args["expansion_epochs"]
 )
 
 self.epochs = self.args["expansion_epochs"]
 self.state = "expansion"
 if len(self._multiple_gpus) > 1:
 network = self._network.module
 else:
 network = self._network
 for p in network.biases.parameters():
 p.requires_grad = False
 self._expansion(train_loader, test_loader, optimizer, scheduler)
 
 
 
 for p in self._network_module_ptr.forward_prototypes.parameters():
 p.requires_grad = False
 for p in self._network_module_ptr.backward_prototypes.parameters():
 p.requires_grad = False
 for p in self._network_module_ptr.new_fc.parameters():
 p.requires_grad = False
 for p in self._network_module_ptr.convnets[-1].parameters():
 p.requires_grad = False
 for p in self._network.biases.parameters():
 p.requires_grad = True
 self.state = "fusion"
 self.epochs = self.args["fusion_epochs"]
 self.per_cls_weights = torch.ones(self._total_classes).to(self._device)
 optimizer = optim.SGD(
 filter(lambda p: p.requires_grad, self._network.parameters()),
 lr=0.05,
 momentum=0.9,
 weight_decay=self.args["weight_decay"],
 )
 for n, p in self._network.named_parameters():
 if p.requires_grad == True:
 print(n)
 scheduler = optim.lr_scheduler.CosineAnnealingLR(
 optimizer=optimizer, T_max=self.args["fusion_epochs"]
 )
 self._fusion(val_loader,test_loader,optimizer,scheduler)

 def _init_train(self, train_loader, test_loader, optimizer, scheduler):
 prog_bar = tqdm(range(self.epochs))
 for _, epoch in enumerate(prog_bar):
 self.train()
 losses = 0.0
 losses_en = 0.0
 correct, total = 0, 0
 for i, (_, inputs, targets) in enumerate(train_loader):
 inputs, targets = inputs.to(
 self._device, non_blocking=True
 ), targets.to(self._device, non_blocking=True)
 logits = self._network(inputs)["logits"]
 loss_en = self.args["energy_weight"] * self.get_energy_loss(inputs,targets,targets)
 loss = F.cross_entropy(logits, targets)
 loss = loss + loss_en
 optimizer.zero_grad()
 loss.backward()
 optimizer.step()
 losses += loss.item()
 losses_en += loss_en.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)
 
 if epoch % 5 == 0:
 test_acc = self._compute_accuracy(self._network, test_loader)
 info = "Task {}, Epoch {}/{} => Loss {:.3f}, Loss_en {:.3f}, Train_accy {:.2f}, Test_accy {:.2f}".format(
 self._cur_task,
 epoch + 1,
 self.args["init_epochs"],
 losses / len(train_loader),
 losses_en / len(train_loader),
 train_acc,
 test_acc,
 )
 else:
 info = "Task {}, Epoch {}/{} => Loss {:.3f}, Loss_en {:.3f}, Train_accy {:.2f}".format(
 self._cur_task,
 epoch + 1,
 self.args["init_epochs"],
 losses / len(train_loader),
 losses_en / len(train_loader),
 train_acc,
 )

 prog_bar.set_description(info)
 logging.info(info)

 def _expansion(self, train_loader, test_loader, optimizer, scheduler):
 prog_bar = tqdm(range(self.epochs))
 for _, epoch in enumerate(prog_bar):
 self.train()
 losses = 0.0
 losses_clf = 0.0
 losses_fe = 0.0
 losses_en = 0.0
 correct, total = 0, 0
 for i, (_, inputs, targets) in enumerate(train_loader):
 inputs, targets = inputs.to(
 self._device, non_blocking=True
 ), targets.to(self._device, non_blocking=True)
 targets = targets.float()
 outputs = self._network(inputs)
 logits,train_logits = (
 outputs["logits"],
 outputs["train_logits"]
 )
 pseudo_targets = targets.clone()
 for task_id in range(self._cur_task+1):
 if task_id == 0:
 pseudo_targets = torch.where(targets<self.data_manager.get_accumulate_tasksize(task_id),torch.Tensor([task_id]).float().to(self._device),pseudo_targets)
 elif task_id == self._cur_task: 
 pseudo_targets = torch.where(targets-self._known_classes+1>0,targets-self._known_classes+task_id,pseudo_targets)
 else:
 pseudo_targets = torch.where((targets<self.data_manager.get_accumulate_tasksize(task_id)) & (targets>self.data_manager.get_accumulate_tasksize(task_id-1)-1),task_id,pseudo_targets)
 
 train_logits[:, list(range(self._cur_task))] /= self.logits_alignment
 loss_clf = F.cross_entropy(train_logits.float(), pseudo_targets)
 loss_fe = torch.tensor(0.).cuda()
 loss_en = self.args["energy_weight"] * self.get_energy_loss(inputs,targets,pseudo_targets)
 loss = loss_clf + loss_fe + loss_en
 optimizer.zero_grad()
 loss.backward()
 optimizer.step()
 losses += loss.item()
 losses_fe += loss_fe.item()
 losses_clf += loss_clf.item()
 losses_en += loss_en.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)
 if epoch % 5 == 0:
 test_acc = self._compute_accuracy(self._network, test_loader)
 info = "Task {}, Epoch {}/{} => Loss {:.3f}, Loss_clf {:.3f}, Loss_fe {:.3f}, Loss_en {:.3f}, Train_accy {:.2f}, Test_accy {:.2f}".format(
 self._cur_task,
 epoch + 1,
 self.epochs,
 losses / len(train_loader),
 losses_clf / len(train_loader),
 losses_fe / len(train_loader),
 losses_en / len(train_loader),
 train_acc,
 test_acc,
 )
 else:
 info = "Task {}, Epoch {}/{} => Loss {:.3f}, Loss_clf {:.3f}, Loss_fe {:.3f}, Loss_en {:.3f}, Train_accy {:.2f}".format(
 self._cur_task,
 epoch + 1,
 self.epochs,
 losses / len(train_loader),
 losses_clf / len(train_loader),
 losses_fe / len(train_loader),
 losses_en / len(train_loader),
 train_acc,
 )
 prog_bar.set_description(info)
 logging.info(info)
 
 def _fusion(self, train_loader, test_loader, optimizer, scheduler):
 prog_bar = tqdm(range(self.epochs))
 for _, epoch in enumerate(prog_bar):
 self.train()
 # self.
 losses = 0.0
 losses_clf = 0.0
 losses_fe = 0.0
 losses_kd = 0.0
 correct, total = 0, 0
 for i, (_, inputs, targets) in enumerate(train_loader):
 inputs, targets = inputs.to(
 self._device, non_blocking=True
 ), targets.to(self._device, non_blocking=True)
 outputs = self._network(inputs)
 logits,train_logits = (
 outputs["logits"],
 outputs["train_logits"]
 )
 
 loss_clf = F.cross_entropy(logits,targets) 
 loss_fe = torch.tensor(0.).cuda()
 loss_kd = torch.tensor(0.).cuda() 
 loss = loss_clf + loss_fe + loss_kd
 optimizer.zero_grad()
 loss.backward()
 optimizer.step()
 losses += loss.item()
 losses_fe += loss_fe.item()
 losses_clf += loss_clf.item()
 losses_kd += (
 self._known_classes / self._total_classes
 ) * loss_kd.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)
 if epoch % 5 == 0:
 test_acc = self._compute_accuracy(self._network, test_loader)
 info = "Task {}, Epoch {}/{} => Loss {:.3f}, Loss_clf {:.3f}, Loss_fe {:.3f}, Loss_kd {:.3f}, Train_accy {:.2f}, Test_accy {:.2f}".format(
 self._cur_task,
 epoch + 1,
 self.epochs,
 losses / len(train_loader),
 losses_clf / len(train_loader),
 losses_fe / len(train_loader),
 losses_kd / len(train_loader),
 train_acc,
 test_acc,
 )
 else:
 info = "Task {}, Epoch {}/{} => Loss {:.3f}, Loss_clf {:.3f}, Loss_fe {:.3f}, Loss_kd {:.3f}, Train_accy {:.2f}".format(
 self._cur_task,
 epoch + 1,
 self.epochs,
 losses / len(train_loader),
 losses_clf / len(train_loader),
 losses_fe / len(train_loader),
 losses_kd / len(train_loader),
 train_acc,
 )
 prog_bar.set_description(info)
 logging.info(info)


 @property
 def samples_old_class(self):
 if self._fixed_memory:
 return self._memory_per_class
 else:
 assert self._total_classes != 0, "Total classes is 0"
 return self._memory_size // self._known_classes

 def samples_new_class(self, index):
 if self.args["dataset"] == "cifar100":
 return 500
 else:
 return self.data_manager.getlen(index)

 def BKD(self, pred, soft, T):
 pred = torch.log_softmax(pred / T, dim=1)
 soft = torch.softmax(soft / T, dim=1)
 soft = soft * self.per_cls_weights
 soft = soft / soft.sum(1)[:, None]
 return -1 * torch.mul(soft, pred).sum() / pred.shape[0]


 def get_energy_loss(self,inputs,targets,pseudo_targets):
 inputs = self.sample_q(inputs)
 
 out = self._network(inputs)
 if self._cur_task == 0:
 targets = targets + self._total_classes
 train_logits, energy_logits = out["logits"], out["energy_logits"]
 else:
 targets = targets + (self._total_classes - self._known_classes) + self._cur_task
 train_logits, energy_logits = out["train_logits"], out["energy_logits"]
 
 logits = torch.cat([train_logits,energy_logits],dim=1)
 
 logits[:,pseudo_targets] = 1e-9 
 energy_loss = F.cross_entropy(logits,targets)
 return energy_loss

 def sample_q(self, replay_buffer, n_steps=3):
 """this func takes in replay_buffer now so we have the option to sample from
 scratch (i.e. replay_buffer==[]). See test_wrn_ebm.py for example.
 """
 self._network_copy = self._network_module_ptr.copy().freeze()
 init_sample = replay_buffer
 init_sample = torch.rot90(init_sample, 2, (2, 3))
 embedding_k = init_sample.clone().detach().requires_grad_(True)
 optimizer_gen = torch.optim.SGD(
 [embedding_k], lr=1e-2)
 for k in range(1, n_steps + 1):
 out = self._network_copy(embedding_k)
 if self._cur_task == 0:
 energy_logits, train_logits = out["energy_logits"], out["logits"]
 else:
 energy_logits, train_logits = out["energy_logits"], out["train_logits"]
 num_forwards = energy_logits.shape[1]
 logits = torch.cat([train_logits,energy_logits],dim=1)
 negative_energy = torch.log(torch.sum(torch.softmax(logits,dim=1)[:,-num_forwards:]))
 optimizer_gen.zero_grad()
 negative_energy.sum().backward()
 optimizer_gen.step()
 embedding_k.data += 1e-3 * \
 torch.randn_like(embedding_k)
 final_samples = embedding_k.detach()
 return final_samples
 
 
 def build_rehearsal_memory_imbalance(self, data_manager, per_class):
 if self._fixed_memory:
 self._construct_exemplar_unified_imbalance(data_manager, per_class,self.random,self.imbalance)
 else:
 self._reduce_exemplar_imbalance(data_manager, per_class,self.random,self.imbalance)
 self._construct_exemplar_imbalance(data_manager, per_class,self.random,self.imbalance)
 
 
 def _reduce_exemplar_imbalance(self, data_manager, m,random,imbalance):
 logging.info('Reducing exemplars...({} per classes)'.format(m))
 dummy_data, dummy_targets = copy.deepcopy(self._data_memory), copy.deepcopy(self._targets_memory)
 self._class_means = np.zeros((self._total_classes, self.feature_dim))
 self._data_memory, self._targets_memory = np.array([]), np.array([])

 for class_idx in range(self._known_classes):
 mask = np.where(dummy_targets == class_idx)[0]
 l = sum(mask)
 if l == 0:
 continue
 if random or imbalance is not None:
 dd, dt = dummy_data[mask][:-1], dummy_targets[mask][:-1]
 else:
 dd, dt = dummy_data[mask][:m], dummy_targets[mask][:m]
 self._data_memory = np.concatenate((self._data_memory, dd)) if len(self._data_memory) != 0 else dd
 self._targets_memory = np.concatenate((self._targets_memory, dt)) if len(self._targets_memory) != 0 else dt

 # Exemplar mean
 idx_dataset = data_manager.get_dataset([], source='train', mode='test', appendent=(dd, dt))
 idx_loader = DataLoader(idx_dataset, batch_size=self.args["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
 mean = np.mean(vectors, axis=0)
 mean = mean / np.linalg.norm(mean)

 self._class_means[class_idx, :] = mean

 def _construct_exemplar_imbalance(self, data_manager, m, random=False,imbalance=None):
 increment = self._total_classes - self._known_classes

 if random:
 '''
 uniform random type
 '''
 selected_exemplars = []
 selected_targets = []
 logging.info("Contructing exmplars, totally random...({} total instances {} classes)".format(increment*m, increment))
 data, targets, idx_dataset = data_manager.get_dataset(np.arange(self._known_classes,self._total_classes),source="train",mode="test",ret_data=True)
 selected_indices = np.random.choice(list(range(len(data))),m*increment,repladce=False)
 for idx in selected_indices:
 selected_exemplars.append(data[idx])
 selected_targets.append(targets[idx])
 selected_exemplars = np.array(selected_exemplars)[:m*increment] 
 selected_targets = np.array(selected_targets)[:m*increment]
 self._data_memory = np.concatenate((self._data_memory, selected_exemplars)) if len(self._data_memory) != 0 \
 else selected_exemplars
 self._targets_memory = np.concatenate((self._targets_memory, selected_targets)) if \
 len(self._targets_memory) != 0 else selected_targets
 else:
 if imbalance is None:
 logging.info('Constructing exemplars...({} per classes)'.format(m))
 ms = np.ones(increment,dtype=int)*m
 elif imbalance>=1:
 '''
 half-half type
 '''
 ms=[m for _ in range(increment)]
 for i in range(increment//2):
 ms[i]-=m//imbalance
 for i in range(increment//2,increment):
 ms[i]+=m//imbalance
 np.random.shuffle(ms)
 ms = np.array(ms,dtype=int)
 logging.info("Constructing exmplars, Imbalance...({} or {} per classes)".format(m-m//imbalance,(m+m//imbalance)))
 elif imbalance<1: 
 '''
 exp type
 '''
 ms = np.array([imbalance**i for i in range(increment)])
 ms = ms/ms.sum()
 tot = m*increment
 ms = (tot*ms).astype(int)
 np.random.shuffle(ms)
 
 else:
 assert 0, "not implemented yet"
 logging.info("ms {}".format(ms))
 for class_idx in range(self._known_classes, self._total_classes):
 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=self.args["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)

 # Select
 selected_exemplars = []
 exemplar_vectors = [] # [n, feature_dim]
 for k in range(1, ms[class_idx-self._known_classes]+1):
 S = np.sum(exemplar_vectors, axis=0) # [feature_dim] sum of selected exemplars vectors
 mu_p = (vectors + S) / k # [n, feature_dim] sum to all vectors
 i = np.argmin(np.sqrt(np.sum((class_mean - mu_p) ** 2, axis=1)))
 selected_exemplars.append(np.array(data[i])) # New object to avoid passing by inference
 exemplar_vectors.append(np.array(vectors[i])) # New object to avoid passing by inference

 vectors = np.delete(vectors, i, axis=0) # Remove it to avoid duplicative selection
 data = np.delete(data, i, axis=0) # Remove it to avoid duplicative selection

 # uniques = np.unique(selected_exemplars, axis=0)
 selected_exemplars = np.array(selected_exemplars)
 if len(selected_exemplars)==0:
 continue
 exemplar_targets = np.full(ms[class_idx-self._known_classes], class_idx)
 self._data_memory = np.concatenate((self._data_memory, selected_exemplars)) if len(self._data_memory) != 0 \
 else selected_exemplars
 self._targets_memory = np.concatenate((self._targets_memory, exemplar_targets)) if \
 len(self._targets_memory) != 0 else exemplar_targets

 # Exemplar mean
 idx_dataset = data_manager.get_dataset([], source='train', mode='test',
 appendent=(selected_exemplars, exemplar_targets))
 idx_loader = DataLoader(idx_dataset, batch_size=self.args["batch_size"], shuffle=False, num_workers=4,pin_memory=True)
 vectors, _ = self._extract_vectors(idx_loader)
 vectors = (vectors.T / (np.linalg.norm(vectors.T, axis=0) + EPSILON)).T
 mean = np.mean(vectors, axis=0)
 mean = mean / np.linalg.norm(mean)

 self._class_means[class_idx, :] = mean
 # self._class_means[class_idx, :] = class_mean

 def _construct_exemplar_unified_imbalance(self, data_manager, m,random,imbalance):
 logging.info('Constructing exemplars for new classes...({} per classes)'.format(m))
 _class_means = np.zeros((self._total_classes, self.feature_dim))
 increment = self._total_classes - self._known_classes

 # Calculate the means of old classes with newly trained network
 for class_idx in range(self._known_classes):
 mask = np.where(self._targets_memory == class_idx)[0]
 if sum(mask) == 0: continue
 class_data, class_targets = self._data_memory[mask], self._targets_memory[mask]

 class_dset = data_manager.get_dataset([], source='train', mode='test',
 appendent=(class_data, class_targets))
 class_loader = DataLoader(class_dset, batch_size=self.args["batch_size"], shuffle=False, num_workers=4)
 vectors, _ = self._extract_vectors(class_loader)
 vectors = (vectors.T / (np.linalg.norm(vectors.T, axis=0) + EPSILON)).T
 mean = np.mean(vectors, axis=0)
 mean = mean / np.linalg.norm(mean)

 _class_means[class_idx, :] = mean

 if random:
 '''
 uniform sample type
 '''
 selected_exemplars = []
 selected_targets = []
 logging.info("Contructing exmplars, totally random...({} total instances {} classes)".format(increment*m, increment))
 data, targets, idx_dataset = data_manager.get_dataset(np.arange(self._known_classes,self._total_classes),source="train",mode="test",ret_data=True)
 selected_indices = np.random.choice(list(range(len(data))),m*increment,replace=False)
 for idx in selected_indices:
 selected_exemplars.append(data[idx])
 selected_targets.append(targets[idx])
 selected_exemplars = np.array(selected_exemplars) 
 selected_targets = np.array(selected_targets)
 self._data_memory = np.concatenate((self._data_memory, selected_exemplars)) if len(self._data_memory) != 0 \
 else selected_exemplars
 self._targets_memory = np.concatenate((self._targets_memory, selected_targets)) if \
 len(self._targets_memory) != 0 else selected_targets
 else:
 if imbalance is None:
 logging.info('Constructing exemplars...({} per classes)'.format(m))
 ms = np.ones(increment,dtype=int)*m
 elif imbalance>=1:
 '''
 half-half type
 '''
 ms=[m for _ in range(increment)]
 for i in range(increment//2):
 ms[i]-=m//imbalance
 for i in range(increment//2,increment):
 ms[i]+=m//imbalance
 np.random.shuffle(ms)
 ms = np.array(ms,dtype=int)
 logging.info("Constructing exmplars, Imbalance...({} or {} per classes)".format(m-m//imbalance,(m+m//imbalance)))
 elif imbalance<1: 
 '''
 exp type
 '''
 ms = np.array([imbalance**i for i in range(increment)])
 ms = ms/ms.sum()
 tot = m*increment
 ms = (tot*ms).astype(int)
 np.random.shuffle(ms)
 
 else:
 assert 0, "not implemented yet"
 logging.info("ms {}".format(ms))
 # Construct exemplars for new classes and calculate the means
 for class_idx in range(self._known_classes, self._total_classes):
 data, targets, class_dset = data_manager.get_dataset(np.arange(class_idx, class_idx+1), source='train',
 mode='test', ret_data=True)
 class_loader = DataLoader(class_dset, batch_size=self.args["batch_size"], shuffle=False, num_workers=4,pin_memory=True)

 vectors, _ = self._extract_vectors(class_loader)
 vectors = (vectors.T / (np.linalg.norm(vectors.T, axis=0) + EPSILON)).T
 class_mean = np.mean(vectors, axis=0)

 # Select
 selected_exemplars = []
 exemplar_vectors = []
 for k in range(1, ms[class_idx-self._known_classes]+1):
 S = np.sum(exemplar_vectors, axis=0) # [feature_dim] sum of selected exemplars vectors
 mu_p = (vectors + S) / k # [n, feature_dim] sum to all vectors
 i = np.argmin(np.sqrt(np.sum((class_mean - mu_p) ** 2, axis=1)))

 selected_exemplars.append(np.array(data[i])) # New object to avoid passing by inference
 exemplar_vectors.append(np.array(vectors[i])) # New object to avoid passing by inference

 vectors = np.delete(vectors, i, axis=0) # Remove it to avoid duplicative selection
 data = np.delete(data, i, axis=0) # Remove it to avoid duplicative selection

 selected_exemplars = np.array(selected_exemplars)
 if len(selected_exemplars)==0:
 continue
 exemplar_targets = np.full(ms[class_idx-self._known_classes], class_idx)
 self._data_memory = np.concatenate((self._data_memory, selected_exemplars)) if len(self._data_memory) != 0 \
 else selected_exemplars
 self._targets_memory = np.concatenate((self._targets_memory, exemplar_targets)) if \
 len(self._targets_memory) != 0 else exemplar_targets

 # Exemplar mean
 exemplar_dset = data_manager.get_dataset([], source='train', mode='test',
 appendent=(selected_exemplars, exemplar_targets))
 exemplar_loader = DataLoader(exemplar_dset, batch_size=self.args["batch_size"], shuffle=False, num_workers=4)
 vectors, _ = self._extract_vectors(exemplar_loader)
 vectors = (vectors.T / (np.linalg.norm(vectors.T, axis=0) + EPSILON)).T
 mean = np.mean(vectors, axis=0)
 mean = mean / np.linalg.norm(mean)

 _class_means[class_idx, :] = mean
 # _class_means[class_idx,:] = class_mean

 self._class_means = _class_means