Spaces:
Running
Running
| #!/usr/bin/env python | |
| import os | |
| from collections import OrderedDict | |
| from timm.models.layers import DropPath | |
| import torch | |
| from torch import nn | |
| from torch.nn import MultiheadAttention | |
| import torch.nn.functional as F | |
| import torch.utils.checkpoint as checkpoint | |
| MODEL_PATH = './' | |
| _MODELS = { | |
| "ViT-B/16": os.path.join(MODEL_PATH, "vit_b16.pth"), | |
| "ViT-L/14": os.path.join(MODEL_PATH, "vit_l14.pth"), | |
| "ViT-L/14_336": os.path.join(MODEL_PATH, "vit_l14_336.pth"), | |
| } | |
| class LayerNorm(nn.LayerNorm): | |
| """Subclass torch's LayerNorm to handle fp16.""" | |
| def forward(self, x): | |
| orig_type = x.dtype | |
| ret = super().forward(x.type(torch.float32)) | |
| return ret.type(orig_type) | |
| class QuickGELU(nn.Module): | |
| def forward(self, x): | |
| return x * torch.sigmoid(1.702 * x) | |
| class Local_MHRA(nn.Module): | |
| def __init__(self, d_model, dw_reduction=1.5, pos_kernel_size=3): | |
| super().__init__() | |
| padding = pos_kernel_size // 2 | |
| re_d_model = int(d_model // dw_reduction) | |
| self.pos_embed = nn.Sequential( | |
| nn.BatchNorm3d(d_model), | |
| nn.Conv3d(d_model, re_d_model, kernel_size=1, stride=1, padding=0), | |
| nn.Conv3d(re_d_model, re_d_model, kernel_size=(pos_kernel_size, 1, 1), stride=(1, 1, 1), padding=(padding, 0, 0), groups=re_d_model), | |
| nn.Conv3d(re_d_model, d_model, kernel_size=1, stride=1, padding=0), | |
| ) | |
| # init zero | |
| print('Init zero for Conv in pos_emb') | |
| nn.init.constant_(self.pos_embed[3].weight, 0) | |
| nn.init.constant_(self.pos_embed[3].bias, 0) | |
| def forward(self, x): | |
| return self.pos_embed(x) | |
| class ResidualAttentionBlock(nn.Module): | |
| def __init__( | |
| self, d_model, n_head, attn_mask=None, drop_path=0.0, | |
| dw_reduction=1.5, no_lmhra=False, double_lmhra=True | |
| ): | |
| super().__init__() | |
| self.n_head = n_head | |
| self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity() | |
| print(f'Drop path rate: {drop_path}') | |
| self.no_lmhra = no_lmhra | |
| self.double_lmhra = double_lmhra | |
| print(f'No L_MHRA: {no_lmhra}') | |
| print(f'Double L_MHRA: {double_lmhra}') | |
| if not no_lmhra: | |
| self.lmhra1 = Local_MHRA(d_model, dw_reduction=dw_reduction) | |
| if double_lmhra: | |
| self.lmhra2 = Local_MHRA(d_model, dw_reduction=dw_reduction) | |
| # spatial | |
| self.attn = MultiheadAttention(d_model, n_head) | |
| self.ln_1 = LayerNorm(d_model) | |
| self.mlp = nn.Sequential(OrderedDict([ | |
| ("c_fc", nn.Linear(d_model, d_model * 4)), | |
| ("gelu", QuickGELU()), | |
| ("c_proj", nn.Linear(d_model * 4, d_model)) | |
| ])) | |
| self.ln_2 = LayerNorm(d_model) | |
| self.attn_mask = attn_mask | |
| def attention(self, x): | |
| self.attn_mask = self.attn_mask.to(dtype=x.dtype, device=x.device) if self.attn_mask is not None else None | |
| return self.attn(x, x, x, need_weights=False, attn_mask=self.attn_mask)[0] | |
| def forward(self, x, T=8, use_checkpoint=False): | |
| # x: 1+HW, NT, C | |
| if not self.no_lmhra: | |
| # Local MHRA | |
| tmp_x = x[1:, :, :] | |
| L, NT, C = tmp_x.shape | |
| N = NT // T | |
| H = W = int(L ** 0.5) | |
| tmp_x = tmp_x.view(H, W, N, T, C).permute(2, 4, 3, 0, 1).contiguous() | |
| tmp_x = tmp_x + self.drop_path(self.lmhra1(tmp_x)) | |
| tmp_x = tmp_x.view(N, C, T, L).permute(3, 0, 2, 1).contiguous().view(L, NT, C) | |
| x = torch.cat([x[:1, :, :], tmp_x], dim=0) | |
| # MHSA | |
| if use_checkpoint: | |
| attn_out = checkpoint.checkpoint(self.attention, self.ln_1(x)) | |
| x = x + self.drop_path(attn_out) | |
| else: | |
| x = x + self.drop_path(self.attention(self.ln_1(x))) | |
| # Local MHRA | |
| if not self.no_lmhra and self.double_lmhra: | |
| tmp_x = x[1:, :, :] | |
| tmp_x = tmp_x.view(H, W, N, T, C).permute(2, 4, 3, 0, 1).contiguous() | |
| tmp_x = tmp_x + self.drop_path(self.lmhra2(tmp_x)) | |
| tmp_x = tmp_x.view(N, C, T, L).permute(3, 0, 2, 1).contiguous().view(L, NT, C) | |
| x = torch.cat([x[:1, :, :], tmp_x], dim=0) | |
| # FFN | |
| if use_checkpoint: | |
| mlp_out = checkpoint.checkpoint(self.mlp, self.ln_2(x)) | |
| x = x + self.drop_path(mlp_out) | |
| else: | |
| x = x + self.drop_path(self.mlp(self.ln_2(x))) | |
| return x | |
| class Extractor(nn.Module): | |
| def __init__( | |
| self, d_model, n_head, attn_mask=None, | |
| mlp_factor=4.0, dropout=0.0, drop_path=0.0, | |
| ): | |
| super().__init__() | |
| self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity() | |
| print(f'Drop path rate: {drop_path}') | |
| self.attn = nn.MultiheadAttention(d_model, n_head) | |
| self.ln_1 = nn.LayerNorm(d_model) | |
| d_mlp = round(mlp_factor * d_model) | |
| self.mlp = nn.Sequential(OrderedDict([ | |
| ("c_fc", nn.Linear(d_model, d_mlp)), | |
| ("gelu", QuickGELU()), | |
| ("dropout", nn.Dropout(dropout)), | |
| ("c_proj", nn.Linear(d_mlp, d_model)) | |
| ])) | |
| self.ln_2 = nn.LayerNorm(d_model) | |
| self.ln_3 = nn.LayerNorm(d_model) | |
| self.attn_mask = attn_mask | |
| # zero init | |
| nn.init.xavier_uniform_(self.attn.in_proj_weight) | |
| nn.init.constant_(self.attn.out_proj.weight, 0.) | |
| nn.init.constant_(self.attn.out_proj.bias, 0.) | |
| nn.init.xavier_uniform_(self.mlp[0].weight) | |
| nn.init.constant_(self.mlp[-1].weight, 0.) | |
| nn.init.constant_(self.mlp[-1].bias, 0.) | |
| def attention(self, x, y): | |
| d_model = self.ln_1.weight.size(0) | |
| q = (x @ self.attn.in_proj_weight[:d_model].T) + self.attn.in_proj_bias[:d_model] | |
| k = (y @ self.attn.in_proj_weight[d_model:-d_model].T) + self.attn.in_proj_bias[d_model:-d_model] | |
| v = (y @ self.attn.in_proj_weight[-d_model:].T) + self.attn.in_proj_bias[-d_model:] | |
| Tx, Ty, N = q.size(0), k.size(0), q.size(1) | |
| q = q.view(Tx, N, self.attn.num_heads, self.attn.head_dim).permute(1, 2, 0, 3) | |
| k = k.view(Ty, N, self.attn.num_heads, self.attn.head_dim).permute(1, 2, 0, 3) | |
| v = v.view(Ty, N, self.attn.num_heads, self.attn.head_dim).permute(1, 2, 0, 3) | |
| aff = (q @ k.transpose(-2, -1) / (self.attn.head_dim ** 0.5)) | |
| aff = aff.softmax(dim=-1) | |
| out = aff @ v | |
| out = out.permute(2, 0, 1, 3).flatten(2) | |
| out = self.attn.out_proj(out) | |
| return out | |
| def forward(self, x, y): | |
| x = x + self.drop_path(self.attention(self.ln_1(x), self.ln_3(y))) | |
| x = x + self.drop_path(self.mlp(self.ln_2(x))) | |
| return x | |
| class Transformer(nn.Module): | |
| def __init__( | |
| self, width, layers, heads, attn_mask=None, backbone_drop_path_rate=0., | |
| use_checkpoint=False, checkpoint_num=[0], t_size=8, dw_reduction=2, | |
| no_lmhra=False, double_lmhra=True, | |
| return_list=[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11], | |
| n_layers=12, n_dim=768, n_head=12, mlp_factor=4.0, drop_path_rate=0., | |
| mlp_dropout=[0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5], | |
| cls_dropout=0.5, num_classes=400, | |
| ): | |
| super().__init__() | |
| self.T = t_size | |
| self.return_list = return_list | |
| # backbone | |
| b_dpr = [x.item() for x in torch.linspace(0, backbone_drop_path_rate, layers)] | |
| self.resblocks = nn.ModuleList([ | |
| ResidualAttentionBlock( | |
| width, heads, attn_mask, | |
| drop_path=b_dpr[i], | |
| dw_reduction=dw_reduction, | |
| no_lmhra=no_lmhra, | |
| double_lmhra=double_lmhra, | |
| ) for i in range(layers) | |
| ]) | |
| # checkpoint | |
| self.use_checkpoint = use_checkpoint | |
| self.checkpoint_num = checkpoint_num | |
| self.n_layers = n_layers | |
| print(f'Use checkpoint: {self.use_checkpoint}') | |
| print(f'Checkpoint number: {self.checkpoint_num}') | |
| # global block | |
| assert n_layers == len(return_list) | |
| if n_layers > 0: | |
| self.temporal_cls_token = nn.Parameter(torch.zeros(1, 1, n_dim)) | |
| self.dpe = nn.ModuleList([ | |
| nn.Conv3d(n_dim, n_dim, kernel_size=3, stride=1, padding=1, bias=True, groups=n_dim) | |
| for i in range(n_layers) | |
| ]) | |
| for m in self.dpe: | |
| nn.init.constant_(m.bias, 0.) | |
| dpr = [x.item() for x in torch.linspace(0, drop_path_rate, n_layers)] | |
| self.dec = nn.ModuleList([ | |
| Extractor( | |
| n_dim, n_head, mlp_factor=mlp_factor, | |
| dropout=mlp_dropout[i], drop_path=dpr[i], | |
| ) for i in range(n_layers) | |
| ]) | |
| self.balance = nn.Parameter(torch.zeros((n_dim))) | |
| self.sigmoid = nn.Sigmoid() | |
| # projection | |
| self.proj = nn.Sequential( | |
| nn.LayerNorm(n_dim), | |
| nn.Dropout(cls_dropout), | |
| nn.Linear(n_dim, num_classes), | |
| ) | |
| def forward(self, x): | |
| T_down = self.T | |
| L, NT, C = x.shape | |
| N = NT // T_down | |
| H = W = int((L - 1) ** 0.5) | |
| if self.n_layers > 0: | |
| cls_token = self.temporal_cls_token.repeat(1, N, 1) | |
| j = -1 | |
| for i, resblock in enumerate(self.resblocks): | |
| if self.use_checkpoint and i < self.checkpoint_num[0]: | |
| x = resblock(x, self.T, use_checkpoint=True) | |
| else: | |
| x = resblock(x, T_down) | |
| if i in self.return_list: | |
| j += 1 | |
| tmp_x = x.clone() | |
| tmp_x = tmp_x.view(L, N, T_down, C) | |
| # dpe | |
| _, tmp_feats = tmp_x[:1], tmp_x[1:] | |
| tmp_feats = tmp_feats.permute(1, 3, 2, 0).reshape(N, C, T_down, H, W) | |
| tmp_feats = self.dpe[j](tmp_feats).view(N, C, T_down, L - 1).permute(3, 0, 2, 1).contiguous() | |
| tmp_x[1:] = tmp_x[1:] + tmp_feats | |
| # global block | |
| tmp_x = tmp_x.permute(2, 0, 1, 3).flatten(0, 1) # T * L, N, C | |
| cls_token = self.dec[j](cls_token, tmp_x) | |
| if self.n_layers > 0: | |
| weight = self.sigmoid(self.balance) | |
| residual = x.view(L, N, T_down, C)[0].mean(1) # L, N, T, C | |
| return self.proj((1 - weight) * cls_token[0, :, :] + weight * residual) | |
| else: | |
| residual = x.view(L, N, T_down, C)[0].mean(1) # L, N, T, C | |
| return self.proj(residual) | |
| class VisionTransformer(nn.Module): | |
| def __init__( | |
| self, | |
| # backbone | |
| input_resolution, patch_size, width, layers, heads, output_dim, backbone_drop_path_rate=0., | |
| use_checkpoint=False, checkpoint_num=[0], t_size=8, kernel_size=3, dw_reduction=1.5, | |
| temporal_downsample=True, | |
| no_lmhra=-False, double_lmhra=True, | |
| # global block | |
| return_list=[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11], | |
| n_layers=12, n_dim=768, n_head=12, mlp_factor=4.0, drop_path_rate=0., | |
| mlp_dropout=[0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5], | |
| cls_dropout=0.5, num_classes=400, | |
| ): | |
| super().__init__() | |
| self.input_resolution = input_resolution | |
| self.output_dim = output_dim | |
| padding = (kernel_size - 1) // 2 | |
| if temporal_downsample: | |
| self.conv1 = nn.Conv3d(3, width, (kernel_size, patch_size, patch_size), (2, patch_size, patch_size), (padding, 0, 0), bias=False) | |
| t_size = t_size // 2 | |
| else: | |
| self.conv1 = nn.Conv3d(3, width, (1, patch_size, patch_size), (1, patch_size, patch_size), (0, 0, 0), bias=False) | |
| scale = width ** -0.5 | |
| self.class_embedding = nn.Parameter(scale * torch.randn(width)) | |
| self.positional_embedding = nn.Parameter(scale * torch.randn((input_resolution // patch_size) ** 2 + 1, width)) | |
| self.ln_pre = LayerNorm(width) | |
| self.transformer = Transformer( | |
| width, layers, heads, dw_reduction=dw_reduction, | |
| backbone_drop_path_rate=backbone_drop_path_rate, | |
| use_checkpoint=use_checkpoint, checkpoint_num=checkpoint_num, t_size=t_size, | |
| no_lmhra=no_lmhra, double_lmhra=double_lmhra, | |
| return_list=return_list, n_layers=n_layers, n_dim=n_dim, n_head=n_head, | |
| mlp_factor=mlp_factor, drop_path_rate=drop_path_rate, mlp_dropout=mlp_dropout, | |
| cls_dropout=cls_dropout, num_classes=num_classes, | |
| ) | |
| def forward(self, x): | |
| x = self.conv1(x) # shape = [*, width, grid, grid] | |
| N, C, T, H, W = x.shape | |
| x = x.permute(0, 2, 3, 4, 1).reshape(N * T, H * W, C) | |
| x = torch.cat([self.class_embedding.to(x.dtype) + torch.zeros(x.shape[0], 1, x.shape[-1], dtype=x.dtype, device=x.device), x], dim=1) # shape = [*, grid ** 2 + 1, width] | |
| x = x + self.positional_embedding.to(x.dtype) | |
| x = self.ln_pre(x) | |
| x = x.permute(1, 0, 2) # NLD -> LND | |
| out = self.transformer(x) | |
| return out | |
| def inflate_weight(weight_2d, time_dim, center=True): | |
| print(f'Init center: {center}') | |
| if center: | |
| weight_3d = torch.zeros(*weight_2d.shape) | |
| weight_3d = weight_3d.unsqueeze(2).repeat(1, 1, time_dim, 1, 1) | |
| middle_idx = time_dim // 2 | |
| weight_3d[:, :, middle_idx, :, :] = weight_2d | |
| else: | |
| weight_3d = weight_2d.unsqueeze(2).repeat(1, 1, time_dim, 1, 1) | |
| weight_3d = weight_3d / time_dim | |
| return weight_3d | |
| def load_state_dict(model, state_dict): | |
| state_dict_3d = model.state_dict() | |
| for k in state_dict.keys(): | |
| if state_dict[k].shape != state_dict_3d[k].shape: | |
| if len(state_dict_3d[k].shape) <= 2: | |
| print(f'Ignore: {k}') | |
| continue | |
| print(f'Inflate: {k}, {state_dict[k].shape} => {state_dict_3d[k].shape}') | |
| time_dim = state_dict_3d[k].shape[2] | |
| state_dict[k] = inflate_weight(state_dict[k], time_dim) | |
| model.load_state_dict(state_dict, strict=False) | |
| def uniformerv2_b16( | |
| pretrained=True, use_checkpoint=False, checkpoint_num=[0], | |
| t_size=16, dw_reduction=1.5, backbone_drop_path_rate=0., | |
| temporal_downsample=True, | |
| no_lmhra=False, double_lmhra=True, | |
| return_list=[8, 9, 10, 11], | |
| n_layers=4, n_dim=768, n_head=12, mlp_factor=4.0, drop_path_rate=0., | |
| mlp_dropout=[0.5, 0.5, 0.5, 0.5], | |
| cls_dropout=0.5, num_classes=400, | |
| ): | |
| model = VisionTransformer( | |
| input_resolution=224, | |
| patch_size=16, | |
| width=768, | |
| layers=12, | |
| heads=12, | |
| output_dim=512, | |
| use_checkpoint=use_checkpoint, | |
| checkpoint_num=checkpoint_num, | |
| t_size=t_size, | |
| dw_reduction=dw_reduction, | |
| backbone_drop_path_rate=backbone_drop_path_rate, | |
| temporal_downsample=temporal_downsample, | |
| no_lmhra=no_lmhra, | |
| double_lmhra=double_lmhra, | |
| return_list=return_list, | |
| n_layers=n_layers, | |
| n_dim=n_dim, | |
| n_head=n_head, | |
| mlp_factor=mlp_factor, | |
| drop_path_rate=drop_path_rate, | |
| mlp_dropout=mlp_dropout, | |
| cls_dropout=cls_dropout, | |
| num_classes=num_classes, | |
| ) | |
| if pretrained: | |
| print('load pretrained weights') | |
| state_dict = torch.load(_MODELS["ViT-B/16"], map_location='cpu') | |
| load_state_dict(model, state_dict) | |
| return model.eval() | |
| def uniformerv2_l14( | |
| pretrained=True, use_checkpoint=False, checkpoint_num=[0], | |
| t_size=16, dw_reduction=1.5, backbone_drop_path_rate=0., | |
| temporal_downsample=True, | |
| no_lmhra=False, double_lmhra=True, | |
| return_list=[20, 21, 22, 23], | |
| n_layers=4, n_dim=1024, n_head=16, mlp_factor=4.0, drop_path_rate=0., | |
| mlp_dropout=[0.5, 0.5, 0.5, 0.5], | |
| cls_dropout=0.5, num_classes=400, | |
| ): | |
| model = VisionTransformer( | |
| input_resolution=224, | |
| patch_size=14, | |
| width=1024, | |
| layers=24, | |
| heads=16, | |
| output_dim=768, | |
| use_checkpoint=use_checkpoint, | |
| checkpoint_num=checkpoint_num, | |
| t_size=t_size, | |
| dw_reduction=dw_reduction, | |
| backbone_drop_path_rate=backbone_drop_path_rate, | |
| temporal_downsample=temporal_downsample, | |
| no_lmhra=no_lmhra, | |
| double_lmhra=double_lmhra, | |
| return_list=return_list, | |
| n_layers=n_layers, | |
| n_dim=n_dim, | |
| n_head=n_head, | |
| mlp_factor=mlp_factor, | |
| drop_path_rate=drop_path_rate, | |
| mlp_dropout=mlp_dropout, | |
| cls_dropout=cls_dropout, | |
| num_classes=num_classes, | |
| ) | |
| if pretrained: | |
| print('load pretrained weights') | |
| state_dict = torch.load(_MODELS["ViT-L/14"], map_location='cpu') | |
| load_state_dict(model, state_dict) | |
| return model.eval() | |
| def uniformerv2_l14_336( | |
| pretrained=True, use_checkpoint=False, checkpoint_num=[0], | |
| t_size=16, dw_reduction=1.5, backbone_drop_path_rate=0., | |
| no_temporal_downsample=True, | |
| no_lmhra=False, double_lmhra=True, | |
| return_list=[20, 21, 22, 23], | |
| n_layers=4, n_dim=1024, n_head=16, mlp_factor=4.0, drop_path_rate=0., | |
| mlp_dropout=[0.5, 0.5, 0.5, 0.5], | |
| cls_dropout=0.5, num_classes=400, | |
| ): | |
| model = VisionTransformer( | |
| input_resolution=336, | |
| patch_size=14, | |
| width=1024, | |
| layers=24, | |
| heads=16, | |
| output_dim=768, | |
| use_checkpoint=use_checkpoint, | |
| checkpoint_num=checkpoint_num, | |
| t_size=t_size, | |
| dw_reduction=dw_reduction, | |
| backbone_drop_path_rate=backbone_drop_path_rate, | |
| no_temporal_downsample=no_temporal_downsample, | |
| no_lmhra=no_lmhra, | |
| double_lmhra=double_lmhra, | |
| return_list=return_list, | |
| n_layers=n_layers, | |
| n_dim=n_dim, | |
| n_head=n_head, | |
| mlp_factor=mlp_factor, | |
| drop_path_rate=drop_path_rate, | |
| mlp_dropout=mlp_dropout, | |
| cls_dropout=cls_dropout, | |
| num_classes=num_classes, | |
| ) | |
| if pretrained: | |
| print('load pretrained weights') | |
| state_dict = torch.load(_MODELS["ViT-L/14_336"], map_location='cpu') | |
| load_state_dict(model, state_dict) | |
| return model.eval() | |
| if __name__ == '__main__': | |
| import time | |
| from fvcore.nn import FlopCountAnalysis | |
| from fvcore.nn import flop_count_table | |
| import numpy as np | |
| seed = 4217 | |
| np.random.seed(seed) | |
| torch.manual_seed(seed) | |
| torch.cuda.manual_seed(seed) | |
| torch.cuda.manual_seed_all(seed) | |
| num_frames = 16 | |
| model = uniformerv2_l14( | |
| pretrained=False, | |
| t_size=num_frames, backbone_drop_path_rate=0., drop_path_rate=0., | |
| dw_reduction=1.5, | |
| no_lmhra=False, | |
| temporal_downsample=True, | |
| return_list=[8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23], | |
| mlp_dropout=[0.5]*16, | |
| n_layers=16 | |
| ) | |
| print(model) | |
| flops = FlopCountAnalysis(model, torch.rand(1, 3, num_frames, 224, 224)) | |
| s = time.time() | |
| print(flop_count_table(flops, max_depth=1)) | |
| print(time.time()-s) |