| |
| """ |
| @inproceedings{DBLP:conf/cvpr/0002ZL0SRSPDP22, |
| author = {Zifeng Wang and |
| Zizhao Zhang and |
| Chen{-}Yu Lee and |
| Han Zhang and |
| Ruoxi Sun and |
| Xiaoqi Ren and |
| Guolong Su and |
| Vincent Perot and |
| Jennifer G. Dy and |
| Tomas Pfister}, |
| title = {Learning to Prompt for Continual Learning}, |
| booktitle = {{IEEE/CVF} Conference on Computer Vision and Pattern Recognition, |
| {CVPR} 2022, New Orleans, LA, USA, June 18-24, 2022}, |
| pages = {139--149}, |
| publisher = {{IEEE}}, |
| year = {2022} |
| } |
| |
| https://arxiv.org/abs/2112.08654 |
| |
| Adapted from https://github.com/GT-RIPL/CODA-Prompt |
| """ |
|
|
| import math |
| import copy |
| import torch |
| import numpy as np |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
| from core.model.backbone.resnet import * |
|
|
| class Model(nn.Module): |
| def __init__(self, backbone, embed_dim, total_cls_num): |
| super().__init__() |
| self.backbone = backbone |
| self.classifier = nn.Linear(embed_dim, total_cls_num, bias=True) |
|
|
| def forward(self, x, train=True): |
| feat, reduce_sim = self.backbone(x, train=train) |
| return self.classifier(feat), reduce_sim |
|
|
| class L2P(nn.Module): |
| def __init__(self, backbone, device, **kwargs): |
| super().__init__() |
|
|
| self.device = device |
| self.init_cls_num = kwargs['init_cls_num'] |
| self.inc_cls_num = kwargs['inc_cls_num'] |
| self.total_cls_num = kwargs['num_class'] |
| self.task_num = kwargs['task_num'] |
| self.embed_dim = kwargs['feat_dim'] |
| self.pull_constraint_coeff = kwargs['pull_constraint_coeff'] |
| self.cur_task_id = 0 |
| self._known_classes = 0 |
| |
| self.network = Model(backbone, self.embed_dim, self.total_cls_num) |
| self.network.backbone.create_prompt( |
| prompt_flag = 'l2p', |
| length = kwargs['prompt_length'], |
| prompt_init = nn.init.uniform_, |
| pool_size = kwargs['pool_size'], |
| top_k = kwargs['top_k'], |
| num_layers = 1, |
| embed_dim = self.embed_dim |
| ) |
| self.network.to(self.device) |
|
|
| self.unfrezeed_params = [] |
| for name, param in self.network.named_parameters(): |
| param.requires_grad_(False) |
| if 'prompt' in name or 'classifier' in name: |
| param.requires_grad_(True) |
| self.unfrezeed_params.append(param) |
|
|
| def before_task(self, task_idx, buffer, train_loader, test_loaders): |
|
|
| self.cur_task_id = task_idx |
|
|
| def after_task(self, task_idx, buffer, train_loader, test_loaders): |
| |
| self._known_classes += self.init_cls_num if task_idx == 0 else self.inc_cls_num |
|
|
| def observe(self, data): |
|
|
| x, y = data['image'].to(self.device), data['label'].to(self.device) |
| logits, reduce_sim = self.network(x, train=True) |
|
|
| if self.cur_task_id == 0: |
| mask = np.arange(self.init_cls_num) |
| else: |
| mask = np.arange(self.inc_cls_num) + self._known_classes |
|
|
| not_mask = np.setdiff1d(np.arange(self.total_cls_num), mask) |
| not_mask = torch.tensor(not_mask, dtype=torch.int64).to(self.device) |
| logits = logits.index_fill(dim=1, index=not_mask, value=float('-inf')) |
|
|
| loss = F.cross_entropy(logits, y) - self.pull_constraint_coeff * reduce_sim |
|
|
| loss.backward() |
| torch.nn.utils.clip_grad_norm_(self.unfrezeed_params, 1.0) |
|
|
| pred = torch.argmax(logits, dim=1) |
| acc = torch.sum(pred == y).item() / x.size(0) |
|
|
| return pred, acc, loss |
|
|
| def inference(self, data): |
| |
| x, y = data['image'].to(self.device), data['label'].to(self.device) |
| logits, _ = self.network(x, train=False) |
|
|
| pred = torch.argmax(logits, dim=1) |
| acc = torch.sum(pred == y).item() / x.size(0) |
| return pred, acc |
|
|
| def get_parameters(self, config): |
|
|
| return self.unfrezeed_params |
|
|
