| | import torch |
| | import torch.nn as nn |
| | import copy |
| |
|
| | from .vit_inflora import VisionTransformer, PatchEmbed, Block, resolve_pretrained_cfg, build_model_with_cfg, checkpoint_filter_fn |
| |
|
| | class ViT_lora_co(VisionTransformer): |
| | def __init__( |
| | self, img_size=224, patch_size=16, in_chans=3, num_classes=1000, global_pool='token', |
| | embed_dim=768, depth=12, num_heads=12, mlp_ratio=4., qkv_bias=True, representation_size=None, |
| | drop_rate=0., attn_drop_rate=0., drop_path_rate=0., weight_init='', init_values=None, |
| | embed_layer=PatchEmbed, norm_layer=None, act_layer=None, block_fn=Block, n_tasks=10, rank=64): |
| |
|
| | super().__init__(img_size=img_size, patch_size=patch_size, in_chans=in_chans, num_classes=num_classes, global_pool=global_pool, |
| | embed_dim=embed_dim, depth=depth, num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, representation_size=representation_size, |
| | drop_rate=drop_rate, attn_drop_rate=attn_drop_rate, drop_path_rate=drop_path_rate, weight_init=weight_init, init_values=init_values, |
| | embed_layer=embed_layer, norm_layer=norm_layer, act_layer=act_layer, block_fn=block_fn, n_tasks=n_tasks, rank=rank) |
| |
|
| | def forward(self, x, task_id, register_blk=-1, get_feat=False, get_cur_feat=False): |
| | x = self.patch_embed(x) |
| | x = torch.cat((self.cls_token.expand(x.shape[0], -1, -1), x), dim=1) |
| |
|
| | x = x + self.pos_embed[:, :x.size(1), :] |
| | x = self.pos_drop(x) |
| |
|
| | prompt_loss = torch.zeros((1,), requires_grad=True).to(x.device) |
| | for i, blk in enumerate(self.blocks): |
| | x = blk(x, task_id, register_blk == i, |
| | get_feat=get_feat, get_cur_feat=get_cur_feat) |
| |
|
| | x = self.norm(x) |
| |
|
| | return x, prompt_loss |
| |
|
| |
|
| | def _create_vision_transformer(variant, pretrained=False, **kwargs): |
| | if kwargs.get('features_only', None): |
| | raise RuntimeError( |
| | 'features_only not implemented for Vision Transformer models.') |
| |
|
| | |
| | |
| | pretrained_cfg = resolve_pretrained_cfg(variant) |
| | default_num_classes = pretrained_cfg['num_classes'] |
| | num_classes = kwargs.get('num_classes', default_num_classes) |
| | repr_size = kwargs.pop('representation_size', None) |
| | if repr_size is not None and num_classes != default_num_classes: |
| | repr_size = None |
| |
|
| | model = build_model_with_cfg( |
| | ViT_lora_co, variant, pretrained, |
| | pretrained_cfg=pretrained_cfg, |
| | representation_size=repr_size, |
| | pretrained_filter_fn=checkpoint_filter_fn, |
| | pretrained_custom_load='npz' in pretrained_cfg['url'], |
| | **kwargs) |
| | return model |
| |
|
| |
|
| | class SiNet_vit(nn.Module): |
| |
|
| | def __init__(self, **args): |
| | ''' |
| | args is a dictionary with the required arguments. |
| | image_encoder is defined in vit_inflora. |
| | class_num is the number of initial class. |
| | ''' |
| | super(SiNet_vit, self).__init__() |
| | model_kwargs = dict(patch_size=16, embed_dim=768, depth=12, |
| | num_heads=12, n_tasks=args["total_sessions"], rank=args["rank"]) |
| | self.image_encoder = _create_vision_transformer( |
| | 'vit_base_patch16_224_in21k', pretrained=True, **model_kwargs) |
| | self.class_num = 1 |
| | self.class_num = args["init_cls"] |
| | self.classifier_pool = nn.ModuleList([ |
| | nn.Linear(args["embd_dim"], self.class_num, bias=True) |
| | for i in range(args["total_sessions"]) |
| | ]) |
| | self.classifier_pool_backup = nn.ModuleList([ |
| | nn.Linear(args["embd_dim"], self.class_num, bias=True) |
| | for i in range(args["total_sessions"]) |
| | ]) |
| | self.numtask = 0 |
| |
|
| | @property |
| | def feature_dim(self): |
| | return self.image_encoder.out_dim |
| |
|
| | def extract_vector(self, image, task=None): |
| | if task == None: |
| | image_features, _ = self.image_encoder(image, self.numtask-1) |
| | else: |
| | image_features, _ = self.image_encoder(image, task) |
| | image_features = image_features[:, 0, :] |
| | return image_features |
| |
|
| | def forward(self, image, get_feat=False, get_cur_feat=False, fc_only=False): |
| | """ |
| | return the output of fully connected layer. |
| | """ |
| | if fc_only: |
| | fc_outs = [] |
| | for ti in range(self.numtask): |
| | fc_out = self.classifier_pool[ti](image) |
| | fc_outs.append(fc_out) |
| | return torch.cat(fc_outs, dim=1) |
| |
|
| | logits = [] |
| | image_features, prompt_loss = self.image_encoder( |
| | image, task_id=self.numtask-1, get_feat=get_feat, get_cur_feat=get_cur_feat) |
| | image_features = image_features[:, 0, :] |
| | image_features = image_features.view(image_features.size(0), -1) |
| | for prompts in [self.classifier_pool[self.numtask-1]]: |
| | logits.append(prompts(image_features)) |
| |
|
| | return { |
| | 'logits': torch.cat(logits, dim=1), |
| | 'features': image_features, |
| | 'prompt_loss': prompt_loss |
| | } |
| |
|
| | def interface(self, image): |
| | image_features, _ = self.image_encoder(image, task_id=self.numtask-1) |
| |
|
| | image_features = image_features[:, 0, :] |
| | image_features = image_features.view(image_features.size(0), -1) |
| |
|
| | logits = [] |
| | for prompt in self.classifier_pool[:self.numtask]: |
| | logits.append(prompt(image_features)) |
| |
|
| | logits = torch.cat(logits, 1) |
| | return logits |
| |
|
| | def update_fc(self, nb_classes): |
| | """ |
| | update the number of tasks. |
| | """ |
| | self.numtask += 1 |
| |
|
| | def classifier_backup(self, task_id): |
| | self.classifier_pool_backup[task_id].load_state_dict( |
| | self.classifier_pool[task_id].state_dict()) |
| |
|
| | def classifier_recall(self): |
| | self.classifier_pool.load_state_dict(self.old_state_dict) |
| |
|
| | def copy(self): |
| | return copy.deepcopy(self) |
| |
|
| | def freeze(self): |
| | for param in self.parameters(): |
| | param.requires_grad = False |
| | self.eval() |
| |
|
| | return self |
| |
|
