Create uniformerv2.py

uniformerv2.py

@@ -0,0 +1,510 @@
+#!/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)