Upload 103 files

models/__init__.py

@@ -0,0 +1,24 @@
+from .yowo.build import build_yowo
+
+
+def build_model(args,
+                d_cfg,
+                m_cfg, 
+                device, 
+                num_classes=3, 
+                trainable=False,
+                resume=None):
+    # build action detector
+    if 'yowo_v2_' in args.version:
+        model, criterion = build_yowo(
+            args=args,
+            d_cfg=d_cfg,
+            m_cfg=m_cfg,
+            device=device,
+            num_classes=num_classes,
+            trainable=trainable,
+            resume=resume
+            )
+
+    return model, criterion
+

models/backbone/__init__.py

@@ -0,0 +1,13 @@
+from .backbone_2d.backbone_2d import Backbone2D
+from .backbone_3d.backbone_3d import Backbone3D
+
+
+def build_backbone_2d(cfg, pretrained=False):
+    backbone = Backbone2D(cfg, pretrained)
+    return backbone, backbone.feat_dims
+
+
+def build_backbone_3d(cfg, pretrained=False):
+    backbone = Backbone3D(cfg, pretrained)
+    return backbone, backbone.feat_dim
+

models/backbone/backbone_2d/backbone_2d.py

@@ -0,0 +1,26 @@
+import torch.nn as nn
+from .cnn_2d import build_2d_cnn
+
+
+class Backbone2D(nn.Module):
+    def __init__(self, cfg, pretrained=False):
+        super().__init__()
+        self.cfg = cfg
+
+        self.backbone, self.feat_dims = build_2d_cnn(cfg, pretrained)
+
+        
+    def forward(self, x):
+        """
+            Input:
+                x: (Tensor) -> [B, C, H, W]
+            Output:
+                y: (List) -> [
+                    (Tensor) -> [B, C1, H1, W1],
+                    (Tensor) -> [B, C2, H2, W2],
+                    (Tensor) -> [B, C3, H3, W3]
+                    ]
+        """
+        feat = self.backbone(x)
+
+        return feat

models/backbone/backbone_2d/cnn_2d/__init__.py

@@ -0,0 +1,18 @@
+# import 2D backbone
+from .yolo_free.yolo_free import build_yolo_free
+
+
+def build_2d_cnn(cfg, pretrained=False):
+    print('==============================')
+    print('2D Backbone: {}'.format(cfg['backbone_2d'].upper()))
+    print('--pretrained: {}'.format(pretrained))
+
+    if cfg['backbone_2d'] in ['yolo_free_nano', 'yolo_free_tiny', \
+                              'yolo_free_large', 'yolo_free_huge']:
+        model, feat_dims = build_yolo_free(cfg['backbone_2d'], pretrained)
+
+    else:
+        print('Unknown 2D Backbone ...')
+        exit()
+
+    return model, feat_dims

models/backbone/backbone_2d/cnn_2d/yolo_free/yolo_free.py

@@ -0,0 +1,222 @@
+import torch
+import numpy as np
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.hub import load_state_dict_from_url
+
+try:
+    from .yolo_free_backbone import build_backbone
+    from .yolo_free_neck import build_neck
+    from .yolo_free_fpn import build_fpn
+    from .yolo_free_head import build_head
+except:
+    from yolo_free_backbone import build_backbone
+    from yolo_free_neck import build_neck
+    from yolo_free_fpn import build_fpn
+    from yolo_free_head import build_head
+
+
+__all__ = ['build_yolo_free']
+
+
+model_urls = {
+    'yolo_free_nano': 'https://github.com/yjh0410/FreeYOLO/releases/download/weight/yolo_free_nano_coco.pth',
+    'yolo_free_tiny': 'https://github.com/yjh0410/FreeYOLO/releases/download/weight/yolo_free_tiny_coco.pth',
+    'yolo_free_large': 'https://github.com/yjh0410/FreeYOLO/releases/download/weight/yolo_free_large_coco.pth',
+}
+
+
+yolo_free_config = {
+    'yolo_free_nano': {
+        # model
+        'backbone': 'shufflenetv2_1.0x',
+        'pretrained': True,
+        'stride': [8, 16, 32],  # P3, P4, P5
+        'anchor_size': None,
+        # neck
+        'neck': 'sppf',
+        'neck_dim': 232,
+        'expand_ratio': 0.5,
+        'pooling_size': 5,
+        'neck_act': 'lrelu',
+        'neck_norm': 'BN',
+        'neck_depthwise': True,
+        # fpn
+        'fpn': 'pafpn_elan',
+        'fpn_size': 'nano',
+        'fpn_dim': [116, 232, 232],
+        'fpn_norm': 'BN',
+        'fpn_act': 'lrelu',
+        'fpn_depthwise': True,
+        # head
+        'head': 'decoupled_head',
+        'head_dim': 64,
+        'head_norm': 'BN',
+        'head_act': 'lrelu',
+        'num_cls_head': 2,
+        'num_reg_head': 2,
+        'head_depthwise': True,
+        },
+
+    'yolo_free_tiny': {
+        # model
+        'backbone': 'elannet_tiny',
+        'pretrained': True,
+        'stride': [8, 16, 32],  # P3, P4, P5
+        # neck
+        'neck': 'spp_block_csp',
+        'neck_dim': 256,
+        'expand_ratio': 0.5,
+        'pooling_size': [5, 9, 13],
+        'neck_act': 'lrelu',
+        'neck_norm': 'BN',
+        'neck_depthwise': False,
+        # fpn
+        'fpn': 'pafpn_elan',
+        'fpn_size': 'tiny', # 'tiny', 'large', 'huge
+        'fpn_dim': [128, 256, 256],
+        'fpn_norm': 'BN',
+        'fpn_act': 'lrelu',
+        'fpn_depthwise': False,
+        # head
+        'head': 'decoupled_head',
+        'head_dim': 64,
+        'head_norm': 'BN',
+        'head_act': 'lrelu',
+        'num_cls_head': 2,
+        'num_reg_head': 2,
+        'head_depthwise': False,
+        },
+
+    'yolo_free_large': {
+        # model
+        'backbone': 'elannet_large',
+        'pretrained': True,
+        'stride': [8, 16, 32],  # P3, P4, P5
+        # neck
+        'neck': 'spp_block_csp',
+        'neck_dim': 512,
+        'expand_ratio': 0.5,
+        'pooling_size': [5, 9, 13],
+        'neck_act': 'silu',
+        'neck_norm': 'BN',
+        'neck_depthwise': False,
+        # fpn
+        'fpn': 'pafpn_elan',
+        'fpn_size': 'large', # 'tiny', 'large', 'huge
+        'fpn_dim': [512, 1024, 512],
+        'fpn_norm': 'BN',
+        'fpn_act': 'silu',
+        'fpn_depthwise': False,
+        # head
+        'head': 'decoupled_head',
+        'head_dim': 256,
+        'head_norm': 'BN',
+        'head_act': 'silu',
+        'num_cls_head': 2,
+        'num_reg_head': 2,
+        'head_depthwise': False,
+        },
+
+}
+
+
+# Anchor-free YOLO
+class FreeYOLO(nn.Module):
+    def __init__(self, cfg):
+        super(FreeYOLO, self).__init__()
+        # --------- Basic Config -----------
+        self.cfg = cfg
+
+        # --------- Network Parameters ----------
+        ## backbone
+        self.backbone, bk_dim = build_backbone(self.cfg['backbone'])
+
+        ## neck
+        self.neck = build_neck(cfg=self.cfg, in_dim=bk_dim[-1], out_dim=self.cfg['neck_dim'])
+        
+        ## fpn
+        self.fpn = build_fpn(cfg=self.cfg, in_dims=self.cfg['fpn_dim'], out_dim=self.cfg['head_dim'])
+
+        ## non-shared heads
+        self.non_shared_heads = nn.ModuleList(
+            [build_head(cfg) 
+            for _ in range(len(cfg['stride']))
+            ])
+
+    def forward(self, x):
+        # backbone
+        feats = self.backbone(x)
+
+        # neck
+        feats['layer4'] = self.neck(feats['layer4'])
+
+        # fpn
+        pyramid_feats = [feats['layer2'], feats['layer3'], feats['layer4']]
+        pyramid_feats = self.fpn(pyramid_feats)
+
+        # non-shared heads
+        all_cls_feats = []
+        all_reg_feats = []
+        for feat, head in zip(pyramid_feats, self.non_shared_heads):
+            # [B, C, H, W]
+            cls_feat, reg_feat = head(feat)
+
+            all_cls_feats.append(cls_feat)
+            all_reg_feats.append(reg_feat)
+
+        return all_cls_feats, all_reg_feats
+
+
+# build FreeYOLO
+def build_yolo_free(model_name='yolo_free_large', pretrained=False):
+    # model config
+    cfg = yolo_free_config[model_name]
+
+    # FreeYOLO
+    model = FreeYOLO(cfg)
+    feat_dims = [model.cfg['head_dim']] * 3
+
+    # Load COCO pretrained weight
+    if pretrained:
+        url = model_urls[model_name]
+
+        # check
+        if url is None:
+            print('No 2D pretrained weight ...')
+            return model, feat_dims
+        else:
+            print('Loading 2D backbone pretrained weight: {}'.format(model_name.upper()))
+
+            # state dict
+            checkpoint = load_state_dict_from_url(url, map_location='cpu')
+            checkpoint_state_dict = checkpoint.pop('model')
+
+            # model state dict
+            model_state_dict = model.state_dict()
+            # check
+            for k in list(checkpoint_state_dict.keys()):
+                if k in model_state_dict:
+                    shape_model = tuple(model_state_dict[k].shape)
+                    shape_checkpoint = tuple(checkpoint_state_dict[k].shape)
+                    if shape_model != shape_checkpoint:
+                        # print(k)
+                        checkpoint_state_dict.pop(k)
+                else:
+                    checkpoint_state_dict.pop(k)
+                    # print(k)
+
+            model.load_state_dict(checkpoint_state_dict, strict=False)
+
+    return model, feat_dims
+
+
+if __name__ == '__main__':
+    model, fpn_dim = build_yolo_free(model_name='yolo_free_nano', pretrained=True)
+    model.eval()
+
+    x = torch.randn(2, 3, 64, 64)
+    cls_feats, reg_feats = model(x)
+
+    for cls_feat, reg_feat in zip(cls_feats, reg_feats):
+        print(cls_feat.shape, reg_feat.shape)

models/backbone/backbone_2d/cnn_2d/yolo_free/yolo_free_backbone.py

@@ -0,0 +1,445 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+__all__ = ['build_backbone']
+
+# ======================  ELAN-Net ==========================
+# ELANNet
+def get_activation(act_type=None):
+    if act_type is None:
+        return nn.Identity()
+    elif act_type == 'relu':
+        return nn.ReLU(inplace=True)
+    elif act_type == 'lrelu':
+        return nn.LeakyReLU(0.1, inplace=True)
+    elif act_type == 'mish':
+        return nn.Mish(inplace=True)
+    elif act_type == 'silu':
+        return nn.SiLU(inplace=True)
+
+
+def get_norm(in_dim, norm_type=None):
+    if norm_type is None:
+        return nn.Identity()
+    elif norm_type == 'BN':
+        return nn.BatchNorm2d(in_dim)
+    elif norm_type == 'GN':
+        return nn.GroupNorm(32, in_dim)
+    elif norm_type == 'IN':
+        return nn.InstanceNorm2d(in_dim)
+
+
+class Conv(nn.Module):
+    def __init__(self, 
+                 c1,                   # in channels
+                 c2,                   # out channels 
+                 k=1,                  # kernel size 
+                 p=0,                  # padding
+                 s=1,                  # padding
+                 d=1,                  # dilation
+                 act_type='silu',
+                 norm_type='BN',       # activation
+                 depthwise=False):
+        super(Conv, self).__init__()
+        convs = []
+        add_bias = False if norm_type else True
+        if depthwise:
+            # depthwise conv
+            convs.append(nn.Conv2d(c1, c1, kernel_size=k, stride=s, padding=p, dilation=d, groups=c1, bias=add_bias))
+            convs.append(get_norm(c1, norm_type))
+            convs.append(get_activation(act_type))
+
+            # pointwise conv
+            convs.append(nn.Conv2d(c1, c2, kernel_size=1, stride=s, padding=0, dilation=d, groups=1, bias=add_bias))
+            convs.append(get_norm(c2, norm_type))
+            convs.append(get_activation(act_type))
+
+        else:
+            convs.append(nn.Conv2d(c1, c2, kernel_size=k, stride=s, padding=p, dilation=d, groups=1, bias=add_bias))
+            convs.append(get_norm(c2, norm_type))
+            convs.append(get_activation(act_type))
+
+        self.convs = nn.Sequential(*convs)
+
+
+    def forward(self, x):
+        return self.convs(x)
+
+
+class ELANBlock(nn.Module):
+    """
+    ELAN BLock of YOLOv7's backbone
+    """
+    def __init__(self, in_dim, out_dim, expand_ratio=0.5, model_size='large', act_type='silu', depthwise=False):
+        super(ELANBlock, self).__init__()
+        inter_dim = int(in_dim * expand_ratio)
+        if model_size == 'tiny':
+            depth = 1
+        elif model_size == 'large':
+            depth = 2
+        elif model_size == 'huge':
+            depth = 3
+        self.cv1 = Conv(in_dim, inter_dim, k=1, act_type=act_type)
+        self.cv2 = Conv(in_dim, inter_dim, k=1, act_type=act_type)
+        self.cv3 = nn.Sequential(*[
+            Conv(inter_dim, inter_dim, k=3, p=1, act_type=act_type, depthwise=depthwise)
+            for _ in range(depth)
+        ])
+        self.cv4 = nn.Sequential(*[
+            Conv(inter_dim, inter_dim, k=3, p=1, act_type=act_type, depthwise=depthwise)
+            for _ in range(depth)
+        ])
+
+        self.out = Conv(inter_dim*4, out_dim, k=1)
+
+
+
+    def forward(self, x):
+        """
+        Input:
+            x: [B, C, H, W]
+        Output:
+            out: [B, 2C, H, W]
+        """
+        x1 = self.cv1(x)
+        x2 = self.cv2(x)
+        x3 = self.cv3(x2)
+        x4 = self.cv4(x3)
+
+        # [B, C, H, W] -> [B, 2C, H, W]
+        out = self.out(torch.cat([x1, x2, x3, x4], dim=1))
+
+        return out
+
+
+class DownSample(nn.Module):
+    def __init__(self, in_dim, act_type='silu', norm_type='BN'):
+        super().__init__()
+        inter_dim = in_dim // 2
+        self.mp = nn.MaxPool2d((2, 2), 2)
+        self.cv1 = Conv(in_dim, inter_dim, k=1, act_type=act_type, norm_type=norm_type)
+        self.cv2 = nn.Sequential(
+            Conv(in_dim, inter_dim, k=1, act_type=act_type, norm_type=norm_type),
+            Conv(inter_dim, inter_dim, k=3, p=1, s=2, act_type=act_type, norm_type=norm_type)
+        )
+
+    def forward(self, x):
+        """
+        Input:
+            x: [B, C, H, W]
+        Output:
+            out: [B, C, H//2, W//2]
+        """
+        # [B, C, H, W] -> [B, C//2, H//2, W//2]
+        x1 = self.cv1(self.mp(x))
+        x2 = self.cv2(x)
+
+        # [B, C, H//2, W//2]
+        out = torch.cat([x1, x2], dim=1)
+
+        return out
+
+
+# ELANNet-Tiny
+class ELANNet_Tiny(nn.Module):
+    """
+    ELAN-Net of YOLOv7-Tiny.
+    """
+    def __init__(self, depthwise=False):
+        super(ELANNet_Tiny, self).__init__()
+        
+        # tiny backbone
+        self.layer_1 = Conv(3, 32, k=3, p=1, s=2, act_type='lrelu', depthwise=depthwise)       # P1/2
+
+        self.layer_2 = nn.Sequential(   
+            Conv(32, 64, k=3, p=1, s=2, act_type='lrelu', depthwise=depthwise),             
+            ELANBlock(in_dim=64, out_dim=64, expand_ratio=0.5,
+                      model_size='tiny', act_type='lrelu', depthwise=depthwise)                  # P2/4
+        )
+        self.layer_3 = nn.Sequential(
+            nn.MaxPool2d((2, 2), 2),             
+            ELANBlock(in_dim=64, out_dim=128, expand_ratio=0.5,
+                      model_size='tiny', act_type='lrelu', depthwise=depthwise)                  # P3/8
+        )
+        self.layer_4 = nn.Sequential(
+            nn.MaxPool2d((2, 2), 2),             
+            ELANBlock(in_dim=128, out_dim=256, expand_ratio=0.5,
+                      model_size='tiny', act_type='lrelu', depthwise=depthwise)                  # P4/16
+        )
+        self.layer_5 = nn.Sequential(
+            nn.MaxPool2d((2, 2), 2),             
+            ELANBlock(in_dim=256, out_dim=512, expand_ratio=0.5,
+                      model_size='tiny', act_type='lrelu', depthwise=depthwise)                  # P5/32
+        )
+
+
+    def forward(self, x):
+        c1 = self.layer_1(x)
+        c2 = self.layer_2(c1)
+        c3 = self.layer_3(c2)
+        c4 = self.layer_4(c3)
+        c5 = self.layer_5(c4)
+
+        outputs = {
+            'layer2': c3,
+            'layer3': c4,
+            'layer4': c5
+        }
+        return outputs
+
+
+# ELANNet-Large
+class ELANNet_Large(nn.Module):
+    """
+    ELAN-Net of YOLOv7.
+    """
+    def __init__(self, depthwise=False):
+        super(ELANNet_Large, self).__init__()
+        
+        # large backbone
+        self.layer_1 = nn.Sequential(
+            Conv(3, 32, k=3, p=1, act_type='silu', depthwise=depthwise),      
+            Conv(32, 64, k=3, p=1, s=2, act_type='silu', depthwise=depthwise),
+            Conv(64, 64, k=3, p=1, act_type='silu', depthwise=depthwise)                                                   # P1/2
+        )
+        self.layer_2 = nn.Sequential(   
+            Conv(64, 128, k=3, p=1, s=2, act_type='silu', depthwise=depthwise),             
+            ELANBlock(in_dim=128, out_dim=256, expand_ratio=0.5,
+                      model_size='large',act_type='silu', depthwise=depthwise)                     # P2/4
+        )
+        self.layer_3 = nn.Sequential(
+            DownSample(in_dim=256, act_type='silu'),             
+            ELANBlock(in_dim=256, out_dim=512, expand_ratio=0.5,
+                      model_size='large',act_type='silu', depthwise=depthwise)                     # P3/8
+        )
+        self.layer_4 = nn.Sequential(
+            DownSample(in_dim=512, act_type='silu'),             
+            ELANBlock(in_dim=512, out_dim=1024, expand_ratio=0.5,
+                      model_size='large',act_type='silu', depthwise=depthwise)                    # P4/16
+        )
+        self.layer_5 = nn.Sequential(
+            DownSample(in_dim=1024, act_type='silu'),             
+            ELANBlock(in_dim=1024, out_dim=1024, expand_ratio=0.25,
+                      model_size='large',act_type='silu', depthwise=depthwise)                  # P5/32
+        )
+
+
+    def forward(self, x):
+        c1 = self.layer_1(x)
+        c2 = self.layer_2(c1)
+        c3 = self.layer_3(c2)
+        c4 = self.layer_4(c3)
+        c5 = self.layer_5(c4)
+
+        outputs = {
+            'layer2': c3,
+            'layer3': c4,
+            'layer4': c5
+        }
+        return outputs
+
+
+## build ELAN-Net
+def build_elannet(model_name='elannet_large'):
+    # model
+    if model_name == 'elannet_large':
+        backbone = ELANNet_Large()
+        feat_dims = [512, 1024, 1024]
+    elif model_name == 'elannet_tiny':
+        backbone = ELANNet_Tiny()
+        feat_dims = [128, 256, 512]
+
+    return backbone, feat_dims
+
+
+# ====================== ShuffleNet-v2 ==========================
+# ShuffleNet-v2
+def channel_shuffle(x, groups):
+    # type: (torch.Tensor, int) -> torch.Tensor
+    batchsize, num_channels, height, width = x.data.size()
+    channels_per_group = num_channels // groups
+
+    # reshape
+    x = x.view(batchsize, groups,
+               channels_per_group, height, width)
+
+    x = torch.transpose(x, 1, 2).contiguous()
+
+    # flatten
+    x = x.view(batchsize, -1, height, width)
+
+    return x
+
+
+class ShuffleV2Block(nn.Module):
+    def __init__(self, inp, oup, stride):
+        super(ShuffleV2Block, self).__init__()
+
+        if not (1 <= stride <= 3):
+            raise ValueError('illegal stride value')
+        self.stride = stride
+
+        branch_features = oup // 2
+        assert (self.stride != 1) or (inp == branch_features << 1)
+
+        if self.stride > 1:
+            self.branch1 = nn.Sequential(
+                self.depthwise_conv(inp, inp, kernel_size=3, stride=self.stride, padding=1),
+                nn.BatchNorm2d(inp),
+                nn.Conv2d(inp, branch_features, kernel_size=1, stride=1, padding=0, bias=False),
+                nn.BatchNorm2d(branch_features),
+                nn.ReLU(inplace=True),
+            )
+        else:
+            self.branch1 = nn.Sequential()
+
+        self.branch2 = nn.Sequential(
+            nn.Conv2d(inp if (self.stride > 1) else branch_features,
+                      branch_features, kernel_size=1, stride=1, padding=0, bias=False),
+            nn.BatchNorm2d(branch_features),
+            nn.ReLU(inplace=True),
+            self.depthwise_conv(branch_features, branch_features, kernel_size=3, stride=self.stride, padding=1),
+            nn.BatchNorm2d(branch_features),
+            nn.Conv2d(branch_features, branch_features, kernel_size=1, stride=1, padding=0, bias=False),
+            nn.BatchNorm2d(branch_features),
+            nn.ReLU(inplace=True),
+        )
+
+    @staticmethod
+    def depthwise_conv(i, o, kernel_size, stride=1, padding=0, bias=False):
+        return nn.Conv2d(i, o, kernel_size, stride, padding, bias=bias, groups=i)
+
+    def forward(self, x):
+        if self.stride == 1:
+            x1, x2 = x.chunk(2, dim=1)
+            out = torch.cat((x1, self.branch2(x2)), dim=1)
+        else:
+            out = torch.cat((self.branch1(x), self.branch2(x)), dim=1)
+
+        out = channel_shuffle(out, 2)
+
+        return out
+
+
+class ShuffleNetV2(nn.Module):
+    def __init__(self,
+                 model_size='1.0x',
+                 out_stages=(2, 3, 4),
+                 with_last_conv=False,
+                 kernal_size=3):
+        super(ShuffleNetV2, self).__init__()
+        print('model size is ', model_size)
+
+        self.stage_repeats = [4, 8, 4]
+        self.model_size = model_size
+        self.out_stages = out_stages
+        self.with_last_conv = with_last_conv
+        self.kernal_size = kernal_size
+        if model_size == '0.5x':
+            self._stage_out_channels = [24, 48, 96, 192]
+        elif model_size == '1.0x':
+            self._stage_out_channels = [24, 116, 232, 464]
+        elif model_size == '1.5x':
+            self._stage_out_channels = [24, 176, 352, 704]
+        elif model_size == '2.0x':
+            self._stage_out_channels = [24, 244, 488, 976]
+        else:
+            raise NotImplementedError
+
+        # building first layer
+        input_channels = 3
+        output_channels = self._stage_out_channels[0]
+        self.conv1 = nn.Sequential(
+            nn.Conv2d(input_channels, output_channels, 3, 2, 1, bias=False),
+            nn.BatchNorm2d(output_channels),
+            nn.ReLU(inplace=True),
+        )
+        input_channels = output_channels
+
+        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+
+        stage_names = ['stage{}'.format(i) for i in [2, 3, 4]]
+        for name, repeats, output_channels in zip(
+                stage_names, self.stage_repeats, self._stage_out_channels[1:]):
+            seq = [ShuffleV2Block(input_channels, output_channels, 2)]
+            for i in range(repeats - 1):
+                seq.append(ShuffleV2Block(output_channels, output_channels, 1))
+            setattr(self, name, nn.Sequential(*seq))
+            input_channels = output_channels
+        
+        self._initialize_weights()
+
+
+    def _initialize_weights(self):
+        print('init weights...')
+        for name, m in self.named_modules():
+            if isinstance(m, nn.Conv2d):
+                if 'first' in name:
+                    nn.init.normal_(m.weight, 0, 0.01)
+                else:
+                    nn.init.normal_(m.weight, 0, 1.0 / m.weight.shape[1])
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0)
+            elif isinstance(m, nn.BatchNorm2d):
+                nn.init.constant_(m.weight, 1)
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0.0001)
+                nn.init.constant_(m.running_mean, 0)
+            elif isinstance(m, nn.BatchNorm1d):
+                nn.init.constant_(m.weight, 1)
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0.0001)
+                nn.init.constant_(m.running_mean, 0)
+            elif isinstance(m, nn.Linear):
+                nn.init.normal_(m.weight, 0, 0.01)
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0)
+
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.maxpool(x)
+        output = {}
+        for i in range(2, 5):
+            stage = getattr(self, 'stage{}'.format(i))
+            x = stage(x)
+            if i in self.out_stages:
+                output['layer{}'.format(i)] = x
+
+        return output
+
+
+## build ShuffleNet-v2
+def build_shufflenetv2(model_size='1.0x'):
+    """Constructs a shufflenetv2 model.
+
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+    """
+    backbone = ShuffleNetV2(model_size=model_size)
+    feat_dims = backbone._stage_out_channels[1:]
+
+    return backbone, feat_dims
+
+
+# build backbone
+def build_backbone(model_name='elannet_large'):
+    if model_name in ['elannet_nano', 'elannet_tiny', 'elannet_large', 'elannet_huge']:
+        return build_elannet(model_name)
+
+    elif model_name in ['shufflenetv2_0.5x', 'shufflenetv2_1.0x']:
+        return build_shufflenetv2(model_size=model_name[-4:])
+        
+
+if __name__ == '__main__':
+    import time
+    model, feats = build_backbone(model_name='shufflenetv2_1.0x')
+    x = torch.randn(1, 3, 224, 224)
+    t0 = time.time()
+    outputs = model(x)
+    t1 = time.time()
+    print('Time: ', t1 - t0)
+    for k in outputs.keys():
+        print(outputs[k].shape)

models/backbone/backbone_2d/cnn_2d/yolo_free/yolo_free_basic.py

@@ -0,0 +1,164 @@
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class SiLU(nn.Module):
+    """export-friendly version of nn.SiLU()"""
+
+    @staticmethod
+    def forward(x):
+        return x * torch.sigmoid(x)
+
+
+def get_conv2d(c1, c2, k, p, s, d, g, padding_mode='ZERO', bias=False):
+    if padding_mode == 'ZERO':
+        conv = nn.Conv2d(c1, c2, k, stride=s, padding=p, dilation=d, groups=g, bias=bias)
+    elif padding_mode == 'SAME':
+        conv = Conv2dSamePadding(c1, c2, k, stride=s, padding=p, dilation=d, groups=g, bias=bias)
+
+    return conv
+
+
+def get_activation(act_type=None):
+    if act_type == 'relu':
+        return nn.ReLU(inplace=True)
+    elif act_type == 'lrelu':
+        return nn.LeakyReLU(0.1, inplace=True)
+    elif act_type == 'mish':
+        return nn.Mish(inplace=True)
+    elif act_type == 'silu':
+        return nn.SiLU(inplace=True)
+
+
+def get_norm(norm_type, dim):
+    if norm_type == 'BN':
+        return nn.BatchNorm2d(dim)
+    elif norm_type == 'GN':
+        return nn.GroupNorm(num_groups=32, num_channels=dim)
+
+
+# Conv2d with "SAME" padding
+class Conv2dSamePadding(nn.Conv2d):
+    """
+    A wrapper around :class:`torch.nn.Conv2d` to support "SAME" padding mode and more features.
+    """
+
+    def __init__(self, *args, **kwargs):
+        """
+        Extra keyword arguments supported in addition to those in `torch.nn.Conv2d`:
+
+        Args:
+            norm (nn.Module, optional): a normalization layer
+            activation (callable(Tensor) -> Tensor): a callable activation function
+
+        It assumes that norm layer is used before activation.
+        """
+
+        # parse padding mode
+        self.padding_method = kwargs.pop("padding", None)
+        if self.padding_method is None:
+            if len(args) >= 5:
+                self.padding_method = args[4]
+            else:
+                self.padding_method = 0  # default padding number
+
+        if isinstance(self.padding_method, str):
+            if self.padding_method.upper() == "SAME":
+                # If the padding mode is `SAME`, it will be manually padded
+                super().__init__(*args, **kwargs, padding=0)
+                # stride
+                if isinstance(self.stride, int):
+                    self.stride = [self.stride] * 2
+                elif len(self.stride) == 1:
+                    self.stride = [self.stride[0]] * 2
+                # kernel size
+                if isinstance(self.kernel_size, int):
+                    self.kernel_size = [self.kernel_size] * 2
+                elif len(self.kernel_size) == 1:
+                    self.kernel_size = [self.kernel_size[0]] * 2
+                # dilation
+                if isinstance(self.dilation, int):
+                    self.dilation = [self.dilation] * 2
+                elif len(self.dilation) == 1:
+                    self.dilation = [self.dilation[0]] * 2
+            else:
+                raise ValueError("Unknown padding method: {}".format(self.padding_method))
+        else:
+            super().__init__(*args, **kwargs, padding=self.padding_method)
+
+    def forward(self, x):
+        if isinstance(self.padding_method, str):
+            if self.padding_method.upper() == "SAME":
+                input_h, input_w = x.shape[-2:]
+                stride_h, stride_w = self.stride
+                kernel_size_h, kernel_size_w = self.kernel_size
+                dilation_h, dilation_w = self.dilation
+
+                output_h = math.ceil(input_h / stride_h)
+                output_w = math.ceil(input_w / stride_w)
+
+                padding_needed_h = max(
+                    0, (output_h - 1) * stride_h + (kernel_size_h - 1) * dilation_h + 1 - input_h
+                )
+                padding_needed_w = max(
+                    0, (output_w - 1) * stride_w + (kernel_size_w - 1) * dilation_w + 1 - input_w
+                )
+
+                left = padding_needed_w // 2
+                right = padding_needed_w - left
+                top = padding_needed_h // 2
+                bottom = padding_needed_h - top
+
+                x = F.pad(x, [left, right, top, bottom])
+            else:
+                raise ValueError("Unknown padding method: {}".format(self.padding_method))
+
+        x = super().forward(x)
+
+        return x
+
+
+# Basic conv layer
+class Conv(nn.Module):
+    def __init__(self, 
+                 c1,                   # in channels
+                 c2,                   # out channels 
+                 k=1,                  # kernel size 
+                 p=0,                  # padding
+                 s=1,                  # padding
+                 d=1,                  # dilation
+                 act_type='',          # activation
+                 norm_type='',         # normalization
+                 padding_mode='ZERO',  # padding mode: "ZERO" or "SAME"
+                 depthwise=False):
+        super(Conv, self).__init__()
+        convs = []
+        add_bias = False if norm_type else True
+        if depthwise:
+            convs.append(get_conv2d(c1, c1, k=k, p=p, s=s, d=d, g=c1, padding_mode=padding_mode, bias=add_bias))
+            # depthwise conv
+            if norm_type:
+                convs.append(get_norm(norm_type, c1))
+            if act_type:
+                convs.append(get_activation(act_type))
+            # pointwise conv
+            convs.append(get_conv2d(c1, c2, k=1, p=0, s=1, d=d, g=1, bias=add_bias))
+            if norm_type:
+                convs.append(get_norm(norm_type, c2))
+            if act_type:
+                convs.append(get_activation(act_type))
+
+        else:
+            convs.append(get_conv2d(c1, c2, k=k, p=p, s=s, d=d, g=1, padding_mode=padding_mode, bias=add_bias))
+            if norm_type:
+                convs.append(get_norm(norm_type, c2))
+            if act_type:
+                convs.append(get_activation(act_type))
+            
+        self.convs = nn.Sequential(*convs)
+
+
+    def forward(self, x):
+        return self.convs(x)

models/backbone/backbone_2d/cnn_2d/yolo_free/yolo_free_fpn.py

@@ -0,0 +1,252 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+try:
+    from yolo_free_basic import Conv
+except:
+    from .yolo_free_basic import Conv
+
+
+class ELANBlock(nn.Module):
+    """
+    ELAN BLock of YOLOv7's head
+    """
+    def __init__(self, in_dim, out_dim, fpn_size='large', depthwise=False, act_type='silu', norm_type='BN'):
+        super(ELANBlock, self).__init__()
+        if fpn_size == 'tiny' or fpn_size =='nano':
+            e1, e2 = 0.25, 1.0
+            width = 2
+            depth = 1
+        elif fpn_size == 'large':
+            e1, e2 = 0.5, 0.5
+            width = 4
+            depth = 1
+        elif fpn_size == 'huge':
+            e1, e2 = 0.5, 0.5
+            width = 4
+            depth = 2
+        inter_dim = int(in_dim * e1)
+        inter_dim2 = int(inter_dim * e2) 
+        self.cv1 = Conv(in_dim, inter_dim, k=1, act_type=act_type, norm_type=norm_type)
+        self.cv2 = Conv(in_dim, inter_dim, k=1, act_type=act_type, norm_type=norm_type)
+        self.cv3 = nn.ModuleList()
+        for idx in range(width):
+            if idx == 0:
+                cvs = [Conv(inter_dim, inter_dim2, k=3, p=1, act_type=act_type, norm_type=norm_type, depthwise=depthwise)]
+            else:
+                cvs = [Conv(inter_dim2, inter_dim2, k=3, p=1, act_type=act_type, norm_type=norm_type, depthwise=depthwise)]
+            # deeper
+            if depth > 1:
+                for _ in range(1, depth):
+                    cvs.append(Conv(inter_dim2, inter_dim2, k=3, p=1, act_type=act_type, norm_type=norm_type, depthwise=depthwise))
+                self.cv3.append(nn.Sequential(*cvs))
+            else:
+                self.cv3.append(cvs[0])
+
+        self.out = Conv(inter_dim*2+inter_dim2*len(self.cv3), out_dim, k=1, act_type=act_type, norm_type=norm_type)
+
+
+    def forward(self, x):
+        """
+        Input:
+            x: [B, C_in, H, W]
+        Output:
+            out: [B, C_out, H, W]
+        """
+        x1 = self.cv1(x)
+        x2 = self.cv2(x)
+        inter_outs = [x1, x2]
+        for m in self.cv3:
+            y1 = inter_outs[-1]
+            y2 = m(y1)
+            inter_outs.append(y2)
+
+        # [B, C_in, H, W] -> [B, C_out, H, W]
+        out = self.out(torch.cat(inter_outs, dim=1))
+
+        return out
+
+
+class DownSample(nn.Module):
+    def __init__(self, in_dim, depthwise=False, act_type='silu', norm_type='BN'):
+        super().__init__()
+        inter_dim = in_dim
+        self.mp = nn.MaxPool2d((2, 2), 2)
+        self.cv1 = Conv(in_dim, inter_dim, k=1, act_type=act_type, norm_type=norm_type)
+        self.cv2 = nn.Sequential(
+            Conv(in_dim, inter_dim, k=1, act_type=act_type, norm_type=norm_type),
+            Conv(inter_dim, inter_dim, k=3, p=1, s=2, act_type=act_type, norm_type=norm_type, depthwise=depthwise)
+        )
+
+    def forward(self, x):
+        """
+        Input:
+            x: [B, C, H, W]
+        Output:
+            out: [B, 2C, H//2, W//2]
+        """
+        # [B, C, H, W] -> [B, C//2, H//2, W//2]
+        x1 = self.cv1(self.mp(x))
+        x2 = self.cv2(x)
+
+        # [B, C, H//2, W//2]
+        out = torch.cat([x1, x2], dim=1)
+
+        return out
+
+
+# PaFPN-ELAN
+class PaFPNELAN(nn.Module):
+    def __init__(self, 
+                 in_dims=[512, 1024, 1024],
+                 out_dim=256,
+                 fpn_size='large',
+                 depthwise=False,
+                 norm_type='BN',
+                 act_type='silu'):
+        super(PaFPNELAN, self).__init__()
+        self.in_dims = in_dims
+        self.out_dim = out_dim
+        c3, c4, c5 = in_dims
+        if fpn_size == 'tiny':
+            width = 0.5
+        elif fpn_size == 'nano':
+            assert depthwise
+            width = 0.5
+        elif fpn_size == 'large':
+            width = 1.0
+        elif fpn_size == 'huge':
+            width = 1.25
+
+        # top dwon
+        ## P5 -> P4
+        self.cv1 = Conv(c5, int(256 * width), k=1, norm_type=norm_type, act_type=act_type)
+        self.cv2 = Conv(c4, int(256 * width), k=1, norm_type=norm_type, act_type=act_type)
+        self.head_elan_1 = ELANBlock(in_dim=int(256 * width) + int(256 * width),
+                                     out_dim=int(256 * width),
+                                     fpn_size=fpn_size,
+                                     depthwise=depthwise,
+                                     norm_type=norm_type,
+                                     act_type=act_type)
+
+        # P4 -> P3
+        self.cv3 = Conv(int(256 * width), int(128 * width), k=1, norm_type=norm_type, act_type=act_type)
+        self.cv4 = Conv(c3, int(128 * width), k=1, norm_type=norm_type, act_type=act_type)
+        self.head_elan_2 = ELANBlock(in_dim=int(128 * width) + int(128 * width),
+                                     out_dim=int(128 * width),  # 128
+                                     fpn_size=fpn_size,
+                                     depthwise=depthwise,
+                                     norm_type=norm_type,
+                                     act_type=act_type)
+
+        # bottom up
+        # P3 -> P4
+        if fpn_size == 'large' or fpn_size == 'huge':
+            self.mp1 = DownSample(int(128 * width), act_type=act_type,
+                                  norm_type=norm_type, depthwise=depthwise)
+        elif fpn_size == 'tiny':
+            self.mp1 = Conv(int(128 * width), int(256 * width), k=3, p=1, s=2,
+                                act_type=act_type, norm_type=norm_type, depthwise=depthwise)
+        elif fpn_size == 'nano':
+            self.mp1 = nn.Sequential(
+                nn.MaxPool2d((2, 2), 2),
+                Conv(int(128 * width), int(256 * width), k=1, act_type=act_type, norm_type=norm_type)
+            )
+        self.head_elan_3 = ELANBlock(in_dim=int(256 * width) + int(256 * width),
+                                     out_dim=int(256 * width),  # 256
+                                     fpn_size=fpn_size,
+                                     depthwise=depthwise,
+                                     norm_type=norm_type,
+                                     act_type=act_type)
+
+        # P4 -> P5
+        if fpn_size == 'large' or fpn_size == 'huge':
+            self.mp2 = DownSample(int(256 * width), act_type=act_type,
+                                  norm_type=norm_type, depthwise=depthwise)
+        elif fpn_size == 'tiny':
+            self.mp2 = Conv(int(256 * width), int(512 * width), k=3, p=1, s=2,
+                                act_type=act_type, norm_type=norm_type, depthwise=depthwise)
+        elif fpn_size == 'nano':
+            self.mp2 = nn.Sequential(
+                nn.MaxPool2d((2, 2), 2),
+                Conv(int(256 * width), int(512 * width), k=1, act_type=act_type, norm_type=norm_type)
+            )
+        self.head_elan_4 = ELANBlock(in_dim=int(512 * width) + c5,
+                                     out_dim=int(512 * width),  # 512
+                                     fpn_size=fpn_size,
+                                     depthwise=depthwise,
+                                     norm_type=norm_type,
+                                     act_type=act_type)
+
+        self.head_conv_1 = Conv(int(128 * width), int(256 * width), k=3, p=1,
+                                act_type=act_type, norm_type=norm_type, depthwise=depthwise)
+        self.head_conv_2 = Conv(int(256 * width), int(512 * width), k=3, p=1,
+                                act_type=act_type, norm_type=norm_type, depthwise=depthwise)
+        self.head_conv_3 = Conv(int(512 * width), int(1024 * width), k=3, p=1,
+                                act_type=act_type, norm_type=norm_type, depthwise=depthwise)
+        # output proj layers
+        if self.out_dim is not None:
+            self.out_layers = nn.ModuleList([
+                Conv(in_dim, self.out_dim, k=1,
+                     norm_type=norm_type, act_type=act_type)
+                     for in_dim in [int(256 * width), int(512 * width), int(1024 * width)]
+                     ])
+
+
+    def forward(self, features):
+        c3, c4, c5 = features
+
+        # Top down
+        ## P5 -> P4
+        c6 = self.cv1(c5)
+        c7 = F.interpolate(c6, scale_factor=2.0)
+        c8 = torch.cat([c7, self.cv2(c4)], dim=1)
+        c9 = self.head_elan_1(c8)
+        ## P4 -> P3
+        c10 = self.cv3(c9)
+        c11 = F.interpolate(c10, scale_factor=2.0)
+        c12 = torch.cat([c11, self.cv4(c3)], dim=1)
+        c13 = self.head_elan_2(c12)
+
+        # Bottom up
+        # p3 -> P4
+        c14 = self.mp1(c13)
+        c15 = torch.cat([c14, c9], dim=1)
+        c16 = self.head_elan_3(c15)
+        # P4 -> P5
+        c17 = self.mp2(c16)
+        c18 = torch.cat([c17, c5], dim=1)
+        c19 = self.head_elan_4(c18)
+
+        c20 = self.head_conv_1(c13)
+        c21 = self.head_conv_2(c16)
+        c22 = self.head_conv_3(c19)
+
+        out_feats = [c20, c21, c22] # [P3, P4, P5]
+        
+        # output proj layers
+        if self.out_dim is not None:
+            out_feats_proj = []
+            for feat, layer in zip(out_feats, self.out_layers):
+                out_feats_proj.append(layer(feat))
+            return out_feats_proj
+
+        return out_feats
+
+
+def build_fpn(cfg, in_dims, out_dim):
+    model = cfg['fpn']
+    print('==============================')
+    print('FPN: {}'.format(model))
+    # build neck
+    if model == 'pafpn_elan':
+        fpn_net = PaFPNELAN(in_dims=in_dims,
+                            out_dim=out_dim,
+                            fpn_size=cfg['fpn_size'],
+                            depthwise=cfg['fpn_depthwise'],
+                            norm_type=cfg['fpn_norm'],
+                            act_type=cfg['fpn_act'])
+                                                        
+
+    return fpn_net

models/backbone/backbone_2d/cnn_2d/yolo_free/yolo_free_head.py

@@ -0,0 +1,51 @@
+import torch
+import torch.nn as nn
+
+try:
+    from yolo_free_basic import Conv
+except:
+    from .yolo_free_basic import Conv
+
+
+class DecoupledHead(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+
+        print('==============================')
+        print('Head: Decoupled Head')
+        self.num_cls_head=cfg['num_cls_head']
+        self.num_reg_head=cfg['num_reg_head']
+        self.act_type=cfg['head_act']
+        self.norm_type=cfg['head_norm']
+        self.head_dim = cfg['head_dim']
+
+        self.cls_feats = nn.Sequential(*[Conv(self.head_dim, 
+                                              self.head_dim, 
+                                              k=3, p=1, s=1, 
+                                              act_type=self.act_type, 
+                                              norm_type=self.norm_type,
+                                              depthwise=cfg['head_depthwise']) for _ in range(self.num_cls_head)])
+        self.reg_feats = nn.Sequential(*[Conv(self.head_dim, 
+                                              self.head_dim, 
+                                              k=3, p=1, s=1, 
+                                              act_type=self.act_type, 
+                                              norm_type=self.norm_type,
+                                              depthwise=cfg['head_depthwise']) for _ in range(self.num_reg_head)])
+
+
+    def forward(self, x):
+        """
+            in_feats: (Tensor) [B, C, H, W]
+        """
+        cls_feats = self.cls_feats(x)
+        reg_feats = self.reg_feats(x)
+
+        return cls_feats, reg_feats
+
+
+# build detection head
+def build_head(cfg):
+    head = DecoupledHead(cfg) 
+
+    return head
+    

models/backbone/backbone_2d/cnn_2d/yolo_free/yolo_free_neck.py

@@ -0,0 +1,164 @@
+import torch
+import torch.nn as nn
+
+try:
+    from yolo_free_basic import Conv
+except:
+    from .yolo_free_basic import Conv
+
+
+# Spatial Pyramid Pooling
+class SPP(nn.Module):
+    """
+        Spatial Pyramid Pooling
+    """
+    def __init__(self, in_dim, out_dim, expand_ratio=0.5, pooling_size=[5, 9, 13], norm_type='BN', act_type='relu'):
+        super(SPP, self).__init__()
+        inter_dim = int(in_dim * expand_ratio)
+        self.cv1 = Conv(in_dim, inter_dim, k=1, act_type=act_type, norm_type=norm_type)
+        self.m = nn.ModuleList(
+            [
+                nn.MaxPool2d(kernel_size=k, stride=1, padding=k // 2)
+                for k in pooling_size
+            ]
+        )
+        
+        self.cv2 = Conv(inter_dim*(len(pooling_size) + 1), out_dim, k=1, act_type=act_type, norm_type=norm_type)
+
+    def forward(self, x):
+        x = self.cv1(x)
+        x = torch.cat([x] + [m(x) for m in self.m], dim=1)
+        x = self.cv2(x)
+
+        return x
+
+
+# SPP block with CSP module
+class SPPBlock(nn.Module):
+    """
+        Spatial Pyramid Pooling Block
+    """
+    def __init__(self,
+                 in_dim,
+                 out_dim,
+                 expand_ratio=0.5,
+                 pooling_size=[5, 9, 13],
+                 act_type='lrelu',
+                 norm_type='BN',
+                 depthwise=False
+                 ):
+        super(SPPBlockCSP, self).__init__()
+        inter_dim = int(in_dim * expand_ratio)
+        self.cv1 = Conv(in_dim, inter_dim, k=1, act_type=act_type, norm_type=norm_type)
+        self.cv2 = nn.Sequential(
+            SPP(inter_dim, 
+                inter_dim, 
+                expand_ratio=1.0, 
+                pooling_size=pooling_size, 
+                act_type=act_type, 
+                norm_type=norm_type),
+        )
+        self.cv3 = Conv(inter_dim * 2, out_dim, k=1, act_type=act_type, norm_type=norm_type)
+
+        
+    def forward(self, x):
+        x1 = self.cv1(x)
+        x2 = self.cv2(x)
+        y = self.cv3(torch.cat([x1, x2], dim=1))
+
+        return y
+
+
+# SPP block with CSP module
+class SPPBlockCSP(nn.Module):
+    """
+        CSP Spatial Pyramid Pooling Block
+    """
+    def __init__(self,
+                 in_dim,
+                 out_dim,
+                 expand_ratio=0.5,
+                 pooling_size=[5, 9, 13],
+                 act_type='lrelu',
+                 norm_type='BN',
+                 depthwise=False
+                 ):
+        super(SPPBlockCSP, self).__init__()
+        inter_dim = int(in_dim * expand_ratio)
+        self.cv1 = Conv(in_dim, inter_dim, k=1, act_type=act_type, norm_type=norm_type)
+        self.cv2 = Conv(in_dim, inter_dim, k=1, act_type=act_type, norm_type=norm_type)
+        self.m = nn.Sequential(
+            Conv(inter_dim, inter_dim, k=3, p=1, 
+                 act_type=act_type, norm_type=norm_type, 
+                 depthwise=depthwise),
+            SPP(inter_dim, 
+                inter_dim, 
+                expand_ratio=1.0, 
+                pooling_size=pooling_size, 
+                act_type=act_type, 
+                norm_type=norm_type),
+            Conv(inter_dim, inter_dim, k=3, p=1, 
+                 act_type=act_type, norm_type=norm_type, 
+                 depthwise=depthwise)
+        )
+        self.cv3 = Conv(inter_dim * 2, out_dim, k=1, act_type=act_type, norm_type=norm_type)
+
+        
+    def forward(self, x):
+        x1 = self.cv1(x)
+        x2 = self.cv2(x)
+        x3 = self.m(x2)
+        y = self.cv3(torch.cat([x1, x3], dim=1))
+
+        return y
+
+
+# Spatial Pyramid Pooling - Fast (SPPF) layer for YOLOv5 by Glenn Jocher
+class SPPF(nn.Module):
+    def __init__(self, in_dim, out_dim, k=5):  # equivalent to SPP(k=(5, 9, 13))
+        super().__init__()
+        inter_dim = in_dim // 2  # hidden channels
+        self.cv1 = Conv(in_dim, inter_dim, k=1)
+        self.cv2 = Conv(inter_dim * 4, out_dim, k=1)
+        self.m = nn.MaxPool2d(kernel_size=k, stride=1, padding=k // 2)
+
+    def forward(self, x):
+        x = self.cv1(x)
+        y1 = self.m(x)
+        y2 = self.m(y1)
+
+        return self.cv2(torch.cat((x, y1, y2, self.m(y2)), 1))
+
+
+def build_neck(cfg, in_dim, out_dim):
+    model = cfg['neck']
+    # build neck
+    if model == 'spp_block':
+        neck = SPPBlock(
+            in_dim, out_dim, 
+            expand_ratio=cfg['expand_ratio'], 
+            pooling_size=cfg['pooling_size'],
+            act_type=cfg['neck_act'],
+            norm_type=cfg['neck_norm'],
+            depthwise=cfg['neck_depthwise']
+            )
+            
+    elif model == 'spp_block_csp':
+        neck = SPPBlockCSP(
+            in_dim, out_dim, 
+            expand_ratio=cfg['expand_ratio'], 
+            pooling_size=cfg['pooling_size'],
+            act_type=cfg['neck_act'],
+            norm_type=cfg['neck_norm'],
+            depthwise=cfg['neck_depthwise']
+            )
+
+    elif model == 'sppf':
+        neck = SPPF(in_dim, out_dim, k=cfg['pooling_size'])
+
+
+    return neck
+
+
+if __name__ == '__main__':
+    pass

models/backbone/backbone_3d/backbone_3d.py

@@ -0,0 +1,68 @@
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .cnn_3d import build_3d_cnn
+
+
+class Conv(nn.Module):
+    def __init__(self, in_dim, out_dim, k=3, p=1, s=1, depthwise=False):
+        super().__init__()
+        if depthwise:
+            self.convs = nn.Sequential(
+                nn.Conv2d(in_dim, in_dim, kernel_size=k, padding=p, stride=s, groups=in_dim, bias=False),
+                nn.BatchNorm2d(out_dim),
+                nn.ReLU(inplace=True),
+                nn.Conv2d(in_dim, out_dim, kernel_size=1, groups=in_dim, bias=False),
+                nn.BatchNorm2d(out_dim),
+                nn.ReLU(inplace=True),
+            )
+        else:
+            self.convs = nn.Sequential(
+                nn.Conv2d(in_dim, out_dim, kernel_size=k, padding=p, stride=s, bias=False),
+                nn.BatchNorm2d(out_dim),
+                nn.ReLU(inplace=True)
+            )
+
+    def forward(self, x):
+        return self.convs(x)
+    
+
+class ConvBlocks(nn.Module):
+    def __init__(self, in_dim, out_dim, nblocks=1, depthwise=False):
+        super().__init__()
+        assert in_dim == out_dim
+
+        conv_block = []
+        for _ in range(nblocks):
+            conv_block.append(
+                Conv(in_dim, out_dim, k=3, p=1, s=1, depthwise=depthwise)
+            )
+        self.conv_block = nn.Sequential(*conv_block)
+
+    def forward(self, x):
+        return self.conv_block(x)
+    
+
+class Backbone3D(nn.Module):
+    def __init__(self, cfg, pretrained=False):
+        super().__init__()
+        self.cfg = cfg
+
+        # 3D CNN
+        self.backbone, self.feat_dim = build_3d_cnn(cfg, pretrained)
+        
+       
+    def forward(self, x):
+        """
+            Input:
+                x: (Tensor) -> [B, C, T, H, W]
+            Output:
+                y: (List) -> [
+                    (Tensor) -> [B, C1, H1, W1],
+                    (Tensor) -> [B, C2, H2, W2],
+                    (Tensor) -> [B, C3, H3, W3]
+                    ]
+        """
+        feat = self.backbone(x)
+
+        return feat

models/backbone/backbone_3d/cnn_3d/__init__.py

@@ -0,0 +1,30 @@
+from .resnet import build_resnet_3d
+from .resnext import build_resnext_3d
+from .shufflnetv2 import build_shufflenetv2_3d
+
+
+def build_3d_cnn(cfg, pretrained=False):
+    print('==============================')
+    print('3D Backbone: {}'.format(cfg['backbone_3d'].upper()))
+    print('--pretrained: {}'.format(pretrained))
+
+    if 'resnet' in cfg['backbone_3d']:
+        model, feat_dims = build_resnet_3d(
+            model_name=cfg['backbone_3d'],
+            pretrained=pretrained
+            )
+    elif 'resnext' in cfg['backbone_3d']:
+        model, feat_dims = build_resnext_3d(
+            model_name=cfg['backbone_3d'],
+            pretrained=pretrained
+            )
+    elif 'shufflenetv2' in cfg['backbone_3d']:
+        model, feat_dims = build_shufflenetv2_3d(
+            model_size=cfg['model_size'],
+            pretrained=pretrained
+            )
+    else:
+        print('Unknown Backbone ...')
+        exit()
+
+    return model, feat_dims