| | from functools import partial |
| | from typing import Literal |
| |
|
| | import torch |
| | import torch.nn as nn |
| | from mmdet.registry import MODELS |
| |
|
| | from mmengine.model import BaseModule |
| | from mmengine.logging import MMLogger |
| |
|
| | from ext.sam import ImageEncoderViT |
| | from ext.meta.sam_meta import meta_dict, checkpoint_dict |
| | from utils.load_checkpoint import load_checkpoint_with_prefix |
| |
|
| |
|
| | @MODELS.register_module() |
| | class SAMBackbone(BaseModule): |
| |
|
| | def __init__( |
| | self, |
| | model_name: Literal['vit_h', 'vit_l', 'vit_b'] = 'vit_h', |
| | fix: bool = True, |
| | init_cfg=None, |
| | ): |
| | assert init_cfg is not None and init_cfg['type'] in \ |
| | ['sam_pretrain', 'Pretrained'], f"{init_cfg['type']} is not supported." |
| | pretrained = init_cfg['checkpoint'] |
| | super().__init__(init_cfg=None) |
| | self.init_cfg = init_cfg |
| | self.logger = MMLogger.get_current_instance() |
| |
|
| | backbone_meta = meta_dict[model_name] |
| |
|
| | backbone = ImageEncoderViT( |
| | depth=backbone_meta['encoder_depth'], |
| | embed_dim=backbone_meta['encoder_embed_dim'], |
| | num_heads=backbone_meta['encoder_num_heads'], |
| | patch_size=backbone_meta['vit_patch_size'], |
| | img_size=backbone_meta['image_size'], |
| | global_attn_indexes=backbone_meta['encoder_global_attn_indexes'], |
| | out_chans=backbone_meta['prompt_embed_dim'], |
| | norm_layer=partial(torch.nn.LayerNorm, eps=1e-6), |
| | qkv_bias=True, |
| | use_rel_pos=True, |
| | mlp_ratio=4, |
| | window_size=14, |
| | ) |
| | if self.init_cfg['type'] == 'sam_pretrain': |
| | checkpoint_path = checkpoint_dict[pretrained] |
| | state_dict = load_checkpoint_with_prefix(checkpoint_path, prefix='image_encoder') |
| | backbone.load_state_dict(state_dict, strict=True) |
| |
|
| | self.stem = backbone.patch_embed |
| | self.pos_embed = backbone.pos_embed |
| |
|
| | self.res_layers = [] |
| | last_pos = 0 |
| | for idx, cur_pos in enumerate(backbone_meta['encoder_global_attn_indexes']): |
| | blocks = backbone.blocks[last_pos:cur_pos + 1] |
| | layer_name = f'layer{idx + 1}' |
| | self.add_module(layer_name, nn.Sequential(*blocks)) |
| | self.res_layers.append(layer_name) |
| | last_pos = cur_pos + 1 |
| |
|
| | self.out_proj = backbone.neck |
| |
|
| | if self.init_cfg['type'] == 'Pretrained': |
| | checkpoint_path = pretrained |
| | state_dict = load_checkpoint_with_prefix(checkpoint_path, prefix=self.init_cfg['prefix']) |
| | self.load_state_dict(state_dict, strict=True) |
| |
|
| | self.model_name = model_name |
| | self.fix = fix |
| | self.model_type = 'vit' |
| | self.output_channels = None |
| | self.out_indices = (0, 1, 2, 3) |
| | if self.fix: |
| | self.train(mode=False) |
| | for name, param in self.named_parameters(): |
| | param.requires_grad = False |
| |
|
| | def init_weights(self): |
| | self.logger.info(f"Init Config for {self.model_name}") |
| | self.logger.info(self.init_cfg) |
| |
|
| | def train(self: torch.nn.Module, mode: bool = True) -> torch.nn.Module: |
| | if not isinstance(mode, bool): |
| | raise ValueError("training mode is expected to be boolean") |
| | if self.fix: |
| | super().train(mode=False) |
| | else: |
| | super().train(mode=mode) |
| | return self |
| |
|
| | def forward_func(self, x): |
| | x = self.stem(x) |
| | x = x + self.pos_embed |
| | outs = [] |
| | for i, layer_name in enumerate(self.res_layers): |
| | res_layer = getattr(self, layer_name) |
| | x = res_layer(x) |
| | if i in self.out_indices: |
| | outs.append(x.permute(0, 3, 1, 2).contiguous()) |
| | outs[-1] = self.out_proj(outs[-1]) |
| | return tuple(outs) |
| |
|
| | def forward(self, x): |
| | if self.fix: |
| | with torch.no_grad(): |
| | outs = self.forward_func(x) |
| | else: |
| | outs = self.forward_func(x) |
| | return outs |
| |
|
