| import math |
| import copy |
| import torch |
| import torch.nn as nn |
| from torch.nn import Parameter |
| import torch.nn.functional as F |
| from .finetune import Finetune |
|
|
| class LWF(Finetune): |
| def __init__(self, backbone, feat_dim, num_class, **kwargs): |
| super().__init__(backbone, feat_dim, num_class, **kwargs) |
| self.kwargs = kwargs |
| self.feat_dim = feat_dim |
| self.classifier = nn.Linear(self.feat_dim, kwargs['init_cls_num']) |
| self.old_fc = None |
| self.init_cls_num = kwargs['init_cls_num'] |
| self.inc_cls_num = kwargs['inc_cls_num'] |
| self.known_cls_num = 0 |
| self.total_cls_num = 0 |
| self.old_backbone = None |
|
|
| def freeze(self,nn): |
| for param in nn.parameters(): |
| param.requires_grad = False |
| nn.eval() |
| return nn |
| |
| def update_fc(self): |
| fc = nn.Linear(self.feat_dim, self.total_cls_num).to(self.device) |
| if self.classifier is not None: |
| |
| self.old_fc = self.freeze(copy.deepcopy(self.classifier)) |
| old_out = self.classifier.out_features |
| weight = copy.deepcopy(self.classifier.weight.data) |
| bias = copy.deepcopy(self.classifier.bias.data) |
| fc.weight.data[:old_out] = weight |
| fc.bias.data[:old_out] = bias |
|
|
| |
| self.classifier = fc |
|
|
| def before_task(self, task_idx, buffer, train_loader, test_loaders): |
| self.task_idx = task_idx |
| self.known_cls_num = self.total_cls_num |
| self.total_cls_num = self.init_cls_num + self.task_idx*self.inc_cls_num |
| self.update_fc() |
| self.loss_fn = nn.CrossEntropyLoss() |
| if task_idx != 0: |
| self.old_backbone = self.freeze(copy.deepcopy(self.backbone)).to(self.device) |
|
|
|
|
| def observe(self, data): |
| x, y = data['image'], data['label'] |
| x = x.to(self.device) |
| y = y.to(self.device) |
| logit = self.classifier(self.backbone(x)['features']) |
|
|
| if self.task_idx == 0: |
| loss = self.loss_fn(logit, y) |
| else: |
| fake_targets = y - self.known_cls_num |
| loss_clf = self.loss_fn(logit[:,self.known_cls_num:],fake_targets) |
| loss_kd = self._KD_loss(logit[:,:self.known_cls_num],self.old_fc(self.old_backbone(x)['features']),T=2) |
| lamda = 3 |
| loss = lamda*loss_kd + loss_clf |
|
|
| pred = torch.argmax(logit, dim=1) |
|
|
| acc = torch.sum(pred == y).item() |
| return pred, acc / x.size(0), loss |
|
|
| def after_task(self, task_idx, buffer, train_loader, test_loaders): |
| pass |
|
|
| def _KD_loss(self, pred, soft, T): |
| pred = torch.log_softmax(pred / T, dim=1) |
| soft = torch.softmax(soft / T, dim=1) |
| return -1 * torch.mul(soft, pred).sum() / pred.shape[0] |
| def _cross_entropy(self, pre, logit): |
| loss = None |
| return loss |
|
|
