Spaces:
Configuration error
Configuration error
| # -*- coding: utf-8 -*- | |
| # Copyright (c) Alibaba, Inc. and its affiliates. | |
| import math | |
| import types | |
| import timm | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| class Slice(nn.Module): | |
| def __init__(self, start_index=1): | |
| super(Slice, self).__init__() | |
| self.start_index = start_index | |
| def forward(self, x): | |
| return x[:, self.start_index:] | |
| class AddReadout(nn.Module): | |
| def __init__(self, start_index=1): | |
| super(AddReadout, self).__init__() | |
| self.start_index = start_index | |
| def forward(self, x): | |
| if self.start_index == 2: | |
| readout = (x[:, 0] + x[:, 1]) / 2 | |
| else: | |
| readout = x[:, 0] | |
| return x[:, self.start_index:] + readout.unsqueeze(1) | |
| class ProjectReadout(nn.Module): | |
| def __init__(self, in_features, start_index=1): | |
| super(ProjectReadout, self).__init__() | |
| self.start_index = start_index | |
| self.project = nn.Sequential(nn.Linear(2 * in_features, in_features), | |
| nn.GELU()) | |
| def forward(self, x): | |
| readout = x[:, 0].unsqueeze(1).expand_as(x[:, self.start_index:]) | |
| features = torch.cat((x[:, self.start_index:], readout), -1) | |
| return self.project(features) | |
| class Transpose(nn.Module): | |
| def __init__(self, dim0, dim1): | |
| super(Transpose, self).__init__() | |
| self.dim0 = dim0 | |
| self.dim1 = dim1 | |
| def forward(self, x): | |
| x = x.transpose(self.dim0, self.dim1) | |
| return x | |
| def forward_vit(pretrained, x): | |
| b, c, h, w = x.shape | |
| _ = pretrained.model.forward_flex(x) | |
| layer_1 = pretrained.activations['1'] | |
| layer_2 = pretrained.activations['2'] | |
| layer_3 = pretrained.activations['3'] | |
| layer_4 = pretrained.activations['4'] | |
| layer_1 = pretrained.act_postprocess1[0:2](layer_1) | |
| layer_2 = pretrained.act_postprocess2[0:2](layer_2) | |
| layer_3 = pretrained.act_postprocess3[0:2](layer_3) | |
| layer_4 = pretrained.act_postprocess4[0:2](layer_4) | |
| unflatten = nn.Sequential( | |
| nn.Unflatten( | |
| 2, | |
| torch.Size([ | |
| h // pretrained.model.patch_size[1], | |
| w // pretrained.model.patch_size[0], | |
| ]), | |
| )) | |
| if layer_1.ndim == 3: | |
| layer_1 = unflatten(layer_1) | |
| if layer_2.ndim == 3: | |
| layer_2 = unflatten(layer_2) | |
| if layer_3.ndim == 3: | |
| layer_3 = unflatten(layer_3) | |
| if layer_4.ndim == 3: | |
| layer_4 = unflatten(layer_4) | |
| layer_1 = pretrained.act_postprocess1[3:len(pretrained.act_postprocess1)]( | |
| layer_1) | |
| layer_2 = pretrained.act_postprocess2[3:len(pretrained.act_postprocess2)]( | |
| layer_2) | |
| layer_3 = pretrained.act_postprocess3[3:len(pretrained.act_postprocess3)]( | |
| layer_3) | |
| layer_4 = pretrained.act_postprocess4[3:len(pretrained.act_postprocess4)]( | |
| layer_4) | |
| return layer_1, layer_2, layer_3, layer_4 | |
| def _resize_pos_embed(self, posemb, gs_h, gs_w): | |
| posemb_tok, posemb_grid = ( | |
| posemb[:, :self.start_index], | |
| posemb[0, self.start_index:], | |
| ) | |
| gs_old = int(math.sqrt(len(posemb_grid))) | |
| posemb_grid = posemb_grid.reshape(1, gs_old, gs_old, | |
| -1).permute(0, 3, 1, 2) | |
| posemb_grid = F.interpolate(posemb_grid, | |
| size=(gs_h, gs_w), | |
| mode='bilinear') | |
| posemb_grid = posemb_grid.permute(0, 2, 3, 1).reshape(1, gs_h * gs_w, -1) | |
| posemb = torch.cat([posemb_tok, posemb_grid], dim=1) | |
| return posemb | |
| def forward_flex(self, x): | |
| b, c, h, w = x.shape | |
| pos_embed = self._resize_pos_embed(self.pos_embed, h // self.patch_size[1], | |
| w // self.patch_size[0]) | |
| B = x.shape[0] | |
| if hasattr(self.patch_embed, 'backbone'): | |
| x = self.patch_embed.backbone(x) | |
| if isinstance(x, (list, tuple)): | |
| x = x[ | |
| -1] # last feature if backbone outputs list/tuple of features | |
| x = self.patch_embed.proj(x).flatten(2).transpose(1, 2) | |
| if getattr(self, 'dist_token', None) is not None: | |
| cls_tokens = self.cls_token.expand( | |
| B, -1, -1) # stole cls_tokens impl from Phil Wang, thanks | |
| dist_token = self.dist_token.expand(B, -1, -1) | |
| x = torch.cat((cls_tokens, dist_token, x), dim=1) | |
| else: | |
| cls_tokens = self.cls_token.expand( | |
| B, -1, -1) # stole cls_tokens impl from Phil Wang, thanks | |
| x = torch.cat((cls_tokens, x), dim=1) | |
| x = x + pos_embed | |
| x = self.pos_drop(x) | |
| for blk in self.blocks: | |
| x = blk(x) | |
| x = self.norm(x) | |
| return x | |
| activations = {} | |
| def get_activation(name): | |
| def hook(model, input, output): | |
| activations[name] = output | |
| return hook | |
| def get_readout_oper(vit_features, features, use_readout, start_index=1): | |
| if use_readout == 'ignore': | |
| readout_oper = [Slice(start_index)] * len(features) | |
| elif use_readout == 'add': | |
| readout_oper = [AddReadout(start_index)] * len(features) | |
| elif use_readout == 'project': | |
| readout_oper = [ | |
| ProjectReadout(vit_features, start_index) for out_feat in features | |
| ] | |
| else: | |
| assert ( | |
| False | |
| ), "wrong operation for readout token, use_readout can be 'ignore', 'add', or 'project'" | |
| return readout_oper | |
| def _make_vit_b16_backbone( | |
| model, | |
| features=[96, 192, 384, 768], | |
| size=[384, 384], | |
| hooks=[2, 5, 8, 11], | |
| vit_features=768, | |
| use_readout='ignore', | |
| start_index=1, | |
| ): | |
| pretrained = nn.Module() | |
| pretrained.model = model | |
| pretrained.model.blocks[hooks[0]].register_forward_hook( | |
| get_activation('1')) | |
| pretrained.model.blocks[hooks[1]].register_forward_hook( | |
| get_activation('2')) | |
| pretrained.model.blocks[hooks[2]].register_forward_hook( | |
| get_activation('3')) | |
| pretrained.model.blocks[hooks[3]].register_forward_hook( | |
| get_activation('4')) | |
| pretrained.activations = activations | |
| readout_oper = get_readout_oper(vit_features, features, use_readout, | |
| start_index) | |
| # 32, 48, 136, 384 | |
| pretrained.act_postprocess1 = nn.Sequential( | |
| readout_oper[0], | |
| Transpose(1, 2), | |
| nn.Unflatten(2, torch.Size([size[0] // 16, size[1] // 16])), | |
| nn.Conv2d( | |
| in_channels=vit_features, | |
| out_channels=features[0], | |
| kernel_size=1, | |
| stride=1, | |
| padding=0, | |
| ), | |
| nn.ConvTranspose2d( | |
| in_channels=features[0], | |
| out_channels=features[0], | |
| kernel_size=4, | |
| stride=4, | |
| padding=0, | |
| bias=True, | |
| dilation=1, | |
| groups=1, | |
| ), | |
| ) | |
| pretrained.act_postprocess2 = nn.Sequential( | |
| readout_oper[1], | |
| Transpose(1, 2), | |
| nn.Unflatten(2, torch.Size([size[0] // 16, size[1] // 16])), | |
| nn.Conv2d( | |
| in_channels=vit_features, | |
| out_channels=features[1], | |
| kernel_size=1, | |
| stride=1, | |
| padding=0, | |
| ), | |
| nn.ConvTranspose2d( | |
| in_channels=features[1], | |
| out_channels=features[1], | |
| kernel_size=2, | |
| stride=2, | |
| padding=0, | |
| bias=True, | |
| dilation=1, | |
| groups=1, | |
| ), | |
| ) | |
| pretrained.act_postprocess3 = nn.Sequential( | |
| readout_oper[2], | |
| Transpose(1, 2), | |
| nn.Unflatten(2, torch.Size([size[0] // 16, size[1] // 16])), | |
| nn.Conv2d( | |
| in_channels=vit_features, | |
| out_channels=features[2], | |
| kernel_size=1, | |
| stride=1, | |
| padding=0, | |
| ), | |
| ) | |
| pretrained.act_postprocess4 = nn.Sequential( | |
| readout_oper[3], | |
| Transpose(1, 2), | |
| nn.Unflatten(2, torch.Size([size[0] // 16, size[1] // 16])), | |
| nn.Conv2d( | |
| in_channels=vit_features, | |
| out_channels=features[3], | |
| kernel_size=1, | |
| stride=1, | |
| padding=0, | |
| ), | |
| nn.Conv2d( | |
| in_channels=features[3], | |
| out_channels=features[3], | |
| kernel_size=3, | |
| stride=2, | |
| padding=1, | |
| ), | |
| ) | |
| pretrained.model.start_index = start_index | |
| pretrained.model.patch_size = [16, 16] | |
| # We inject this function into the VisionTransformer instances so that | |
| # we can use it with interpolated position embeddings without modifying the library source. | |
| pretrained.model.forward_flex = types.MethodType(forward_flex, | |
| pretrained.model) | |
| pretrained.model._resize_pos_embed = types.MethodType( | |
| _resize_pos_embed, pretrained.model) | |
| return pretrained | |
| def _make_pretrained_vitl16_384(pretrained, use_readout='ignore', hooks=None): | |
| model = timm.create_model('vit_large_patch16_384', pretrained=pretrained) | |
| hooks = [5, 11, 17, 23] if hooks is None else hooks | |
| return _make_vit_b16_backbone( | |
| model, | |
| features=[256, 512, 1024, 1024], | |
| hooks=hooks, | |
| vit_features=1024, | |
| use_readout=use_readout, | |
| ) | |
| def _make_pretrained_vitb16_384(pretrained, use_readout='ignore', hooks=None): | |
| model = timm.create_model('vit_base_patch16_384', pretrained=pretrained) | |
| hooks = [2, 5, 8, 11] if hooks is None else hooks | |
| return _make_vit_b16_backbone(model, | |
| features=[96, 192, 384, 768], | |
| hooks=hooks, | |
| use_readout=use_readout) | |
| def _make_pretrained_deitb16_384(pretrained, use_readout='ignore', hooks=None): | |
| model = timm.create_model('vit_deit_base_patch16_384', | |
| pretrained=pretrained) | |
| hooks = [2, 5, 8, 11] if hooks is None else hooks | |
| return _make_vit_b16_backbone(model, | |
| features=[96, 192, 384, 768], | |
| hooks=hooks, | |
| use_readout=use_readout) | |
| def _make_pretrained_deitb16_distil_384(pretrained, | |
| use_readout='ignore', | |
| hooks=None): | |
| model = timm.create_model('vit_deit_base_distilled_patch16_384', | |
| pretrained=pretrained) | |
| hooks = [2, 5, 8, 11] if hooks is None else hooks | |
| return _make_vit_b16_backbone( | |
| model, | |
| features=[96, 192, 384, 768], | |
| hooks=hooks, | |
| use_readout=use_readout, | |
| start_index=2, | |
| ) | |
| def _make_vit_b_rn50_backbone( | |
| model, | |
| features=[256, 512, 768, 768], | |
| size=[384, 384], | |
| hooks=[0, 1, 8, 11], | |
| vit_features=768, | |
| use_vit_only=False, | |
| use_readout='ignore', | |
| start_index=1, | |
| ): | |
| pretrained = nn.Module() | |
| pretrained.model = model | |
| if use_vit_only is True: | |
| pretrained.model.blocks[hooks[0]].register_forward_hook( | |
| get_activation('1')) | |
| pretrained.model.blocks[hooks[1]].register_forward_hook( | |
| get_activation('2')) | |
| else: | |
| pretrained.model.patch_embed.backbone.stages[0].register_forward_hook( | |
| get_activation('1')) | |
| pretrained.model.patch_embed.backbone.stages[1].register_forward_hook( | |
| get_activation('2')) | |
| pretrained.model.blocks[hooks[2]].register_forward_hook( | |
| get_activation('3')) | |
| pretrained.model.blocks[hooks[3]].register_forward_hook( | |
| get_activation('4')) | |
| pretrained.activations = activations | |
| readout_oper = get_readout_oper(vit_features, features, use_readout, | |
| start_index) | |
| if use_vit_only is True: | |
| pretrained.act_postprocess1 = nn.Sequential( | |
| readout_oper[0], | |
| Transpose(1, 2), | |
| nn.Unflatten(2, torch.Size([size[0] // 16, size[1] // 16])), | |
| nn.Conv2d( | |
| in_channels=vit_features, | |
| out_channels=features[0], | |
| kernel_size=1, | |
| stride=1, | |
| padding=0, | |
| ), | |
| nn.ConvTranspose2d( | |
| in_channels=features[0], | |
| out_channels=features[0], | |
| kernel_size=4, | |
| stride=4, | |
| padding=0, | |
| bias=True, | |
| dilation=1, | |
| groups=1, | |
| ), | |
| ) | |
| pretrained.act_postprocess2 = nn.Sequential( | |
| readout_oper[1], | |
| Transpose(1, 2), | |
| nn.Unflatten(2, torch.Size([size[0] // 16, size[1] // 16])), | |
| nn.Conv2d( | |
| in_channels=vit_features, | |
| out_channels=features[1], | |
| kernel_size=1, | |
| stride=1, | |
| padding=0, | |
| ), | |
| nn.ConvTranspose2d( | |
| in_channels=features[1], | |
| out_channels=features[1], | |
| kernel_size=2, | |
| stride=2, | |
| padding=0, | |
| bias=True, | |
| dilation=1, | |
| groups=1, | |
| ), | |
| ) | |
| else: | |
| pretrained.act_postprocess1 = nn.Sequential(nn.Identity(), | |
| nn.Identity(), | |
| nn.Identity()) | |
| pretrained.act_postprocess2 = nn.Sequential(nn.Identity(), | |
| nn.Identity(), | |
| nn.Identity()) | |
| pretrained.act_postprocess3 = nn.Sequential( | |
| readout_oper[2], | |
| Transpose(1, 2), | |
| nn.Unflatten(2, torch.Size([size[0] // 16, size[1] // 16])), | |
| nn.Conv2d( | |
| in_channels=vit_features, | |
| out_channels=features[2], | |
| kernel_size=1, | |
| stride=1, | |
| padding=0, | |
| ), | |
| ) | |
| pretrained.act_postprocess4 = nn.Sequential( | |
| readout_oper[3], | |
| Transpose(1, 2), | |
| nn.Unflatten(2, torch.Size([size[0] // 16, size[1] // 16])), | |
| nn.Conv2d( | |
| in_channels=vit_features, | |
| out_channels=features[3], | |
| kernel_size=1, | |
| stride=1, | |
| padding=0, | |
| ), | |
| nn.Conv2d( | |
| in_channels=features[3], | |
| out_channels=features[3], | |
| kernel_size=3, | |
| stride=2, | |
| padding=1, | |
| ), | |
| ) | |
| pretrained.model.start_index = start_index | |
| pretrained.model.patch_size = [16, 16] | |
| # We inject this function into the VisionTransformer instances so that | |
| # we can use it with interpolated position embeddings without modifying the library source. | |
| pretrained.model.forward_flex = types.MethodType(forward_flex, | |
| pretrained.model) | |
| # We inject this function into the VisionTransformer instances so that | |
| # we can use it with interpolated position embeddings without modifying the library source. | |
| pretrained.model._resize_pos_embed = types.MethodType( | |
| _resize_pos_embed, pretrained.model) | |
| return pretrained | |
| def _make_pretrained_vitb_rn50_384(pretrained, | |
| use_readout='ignore', | |
| hooks=None, | |
| use_vit_only=False): | |
| model = timm.create_model('vit_base_resnet50_384', pretrained=pretrained) | |
| # model = timm.create_model('vit_base_r50_s16_384.orig_in21k_ft_in1k', pretrained=pretrained) | |
| hooks = [0, 1, 8, 11] if hooks is None else hooks | |
| return _make_vit_b_rn50_backbone( | |
| model, | |
| features=[256, 512, 768, 768], | |
| size=[384, 384], | |
| hooks=hooks, | |
| use_vit_only=use_vit_only, | |
| use_readout=use_readout, | |
| ) | |