| import torch |
| import copy |
| from models.utils.continual_model import ContinualModel |
| from utils.args import add_management_args, add_experiment_args, add_rehearsal_args, ArgumentParser |
| from utils.buffer import Buffer |
| import torch.nn.functional as F |
| from utils.args import * |
|
|
|
|
|
|
| def add_parser(parser): |
| parser.add_argument('--weighta', type=float, help='Penalty weight for idempotence distillation.') |
| parser.add_argument('--weightb', type=float, help='Penalty weight for current idempotence distillation.') |
| parser.add_argument('--weightc', type=float, help='Penalty weight for er.') |
| parser.add_argument('--weightmask', type=float, help='Penalty weight for mask ratio.') |
| parser.add_argument("--class_balance", type=str2bool, default=True, |
| help="If set, the memory buffer will be balanced by class") |
| return parser |
| def get_parser() -> ArgumentParser: |
| parser = ArgumentParser(description='Idempotent Continual learning via' |
| ' Experience Replay.') |
| add_management_args(parser) |
| add_experiment_args(parser) |
| add_rehearsal_args(parser) |
| parser = add_parser(parser) |
| return parser |
|
|
| class Ider(ContinualModel): |
| NAME = 'ider' |
| COMPATIBILITY = ['class-il', 'domain-il', 'task-il', 'general-continual'] |
|
|
| def __init__(self, backbone, loss, args, transform): |
| super(Ider, self).__init__(backbone, loss, args, transform) |
| self.buffer = Buffer(self.args.buffer_size, self.device,class_balance = self.args.class_balance) |
| self.ft=True |
| self.task=0 |
| self.num_classs = backbone.num_classes |
| self.s = backbone.num_classes |
| self.first_task = True |
| self.old_model=self.deepcopy_model(self.net) |
|
|
| |
| |
| def observe(self, inputs, labels, not_aug_inputs): |
| batch_size, _, H, W = inputs.shape |
|
|
| self.opt.zero_grad() |
| mask_current = torch.rand(1) > self.args.weightmask |
| y_0_current = F.one_hot(labels, self.num_classs).float() if mask_current else torch.ones(batch_size, self.num_classs).to(self.device) /self.s |
| z_current = self.net.f1(inputs) |
| y_1_current, z1_current = self.net.f2(z_current, y_0_current) |
| y_2_current, z2_current = self.net.f2(z_current , y_1_current.softmax(-1)) |
| loss_supervised_1 = self.loss(y_1_current, labels) |
| loss_supervised_2 = self.loss(y_2_current, labels) |
|
|
| loss = 0.5*(loss_supervised_1 + loss_supervised_2) |
| |
| if self.args.weightb!=0 and self.task>0: |
| y_current_mask = torch.ones(batch_size, self.num_classs).to(self.device) /self.s |
| z = self.net.f1(inputs) |
| y_1, z1 = self.net.f2(z, y_current_mask) |
| z_old = self.old_model.f1(inputs) |
| y_2, z2 = self.old_model.f2(z_old, y_1.softmax(-1)) |
| loss += self.args.weightb*F.mse_loss(y_1, y_2) |
| |
| if not self.buffer.is_empty() and self.args.weightc !=0: |
| buf_inputs, buf_labels,_,_,_ = self.buffer.get_data( |
| self.args.minibatch_size, transform=self.transform) |
| batch_size, _, H, W = buf_inputs.shape |
| mask = torch.rand(1) > self.args.weightmask |
| y_0_buf = F.one_hot(buf_labels, self.num_classs).float() if mask else torch.ones(batch_size, self.num_classs).to(self.device) /self.s |
| z_buf = self.net.f1(buf_inputs) |
| y_1_buf, z1_buf = self.net.f2(z_buf, y_0_buf) |
| y_2_buf, z2_buf = self.net.f2(z_buf , y_1_buf.softmax(-1)) |
| loss_supervised_1_buf = self.loss(y_1_buf, buf_labels) |
| loss_supervised_2_buf = self.loss(y_2_buf, buf_labels) |
|
|
| loss += self.args.weightc*(loss_supervised_1_buf + loss_supervised_2_buf) |
|
|
| if not self.buffer.is_empty() and self.task>0 and self.args.weighta!=0: |
| buf_inputs, buf_labels,_,_,_ = self.buffer.get_data( |
| self.args.minibatch_size, transform=self.transform) |
| batch_size, _, H, W = buf_inputs.shape |
| |
| y_0 = torch.ones(batch_size, self.num_classs).to(self.device) /self.s |
| |
| z = self.net.f1(buf_inputs) |
| y_1, z1 = self.net.f2(z, y_0) |
| |
| z_old = self.old_model.f1(buf_inputs) |
| y_2, z2 = self.old_model.f2(z_old, y_1.softmax(-1)) |
| loss_unsupervised_y = F.mse_loss(y_1, y_2) |
| loss += self.args.weighta * loss_unsupervised_y |
| |
| |
| loss.backward() |
| self.opt.step() |
| self.buffer.add_data(examples=not_aug_inputs, labels=labels,logits=y_1_current.data, logits2=y_2_current.data,mask=y_0_current) |
|
|
| return loss.item() |
| def end_task(self, dataset): |
| print('\n\n') |
| self.task+=1 |
| print(self.task) |
|
|
| if self.first_task: |
| self.first_task = False |
| self.old_model = self.deepcopy_model(self.net).to(self.device) |
| else: |
| self.old_model = self.deepcopy_model(self.net).to(self.device) |
| |
| @staticmethod |
| def deepcopy_model(model): |
| model_copy = copy.deepcopy(model) |
| return model_copy |
|
|
