RT-ODLab/models/detectors/rtcdet/rtcdet_backbone.py

import torch
import torch.nn as nn

try:
    from .rtcdet_basic import Conv, RTCBlock
except:
    from rtcdet_basic import Conv, RTCBlock


# Pretrained weights
model_urls = {
    # IN-1k classification pretrained
    "rtcnet_n": "https://github.com/yjh0410/image_classification_pytorch/releases/download/weight/elan_cspnet_nano.pth",
    "rtcnet_s": "https://github.com/yjh0410/image_classification_pytorch/releases/download/weight/elan_cspnet_small.pth",
    "rtcnet_m": None,
    "rtcnet_l": None,
    "rtcnet_x": None,
    # IN-1k MIM pretrained
    "mae_rtcnet_n": None,
    "mae_rtcnet_s": None,
    "mae_rtcnet_m": None,
    "mae_rtcnet_l": None,
    "mae_rtcnet_x": None,
}


# ---------------------------- Basic functions ----------------------------
## Real-time Convolutional Backbone
class RTCBackbone(nn.Module):
    def __init__(self, width=1.0, depth=1.0, ratio=1.0, act_type='silu', norm_type='BN', depthwise=False):
        super(RTCBackbone, self).__init__()
        # ---------------- Basic parameters ----------------
        self.width_factor = width
        self.depth_factor = depth
        self.last_stage_factor = ratio
        self.feat_dims = [round(64 * width), round(128 * width), round(256 * width), round(512 * width), round(512 * width * ratio)]
        # ---------------- Network parameters ----------------
        ## P1/2
        self.layer_1 = Conv(3, self.feat_dims[0], k=3, p=1, s=2, act_type=act_type, norm_type=norm_type)
        ## P2/4
        self.layer_2 = nn.Sequential(
            Conv(self.feat_dims[0], self.feat_dims[1], k=3, p=1, s=2, act_type=act_type, norm_type=norm_type),
            RTCBlock(in_dim     = self.feat_dims[1],
                     out_dim    = self.feat_dims[1],
                     num_blocks = round(3*depth),
                     shortcut   = True,
                     act_type   = act_type,
                     norm_type  = norm_type,
                     depthwise  = depthwise)
        )
        ## P3/8
        self.layer_3 = nn.Sequential(
            Conv(self.feat_dims[1], self.feat_dims[2], k=3, p=1, s=2, act_type=act_type, norm_type=norm_type),
            RTCBlock(in_dim     = self.feat_dims[2],
                     out_dim    = self.feat_dims[2],
                     num_blocks = round(6*depth),
                     shortcut   = True,
                     act_type   = act_type,
                     norm_type  = norm_type,
                     depthwise  = depthwise)
        )
        ## P4/16
        self.layer_4 = nn.Sequential(
            Conv(self.feat_dims[2], self.feat_dims[3], k=3, p=1, s=2, act_type=act_type, norm_type=norm_type),
            RTCBlock(in_dim     = self.feat_dims[3],
                     out_dim    = self.feat_dims[3],
                     num_blocks = round(6*depth),
                     shortcut   = True,
                     act_type   = act_type,
                     norm_type  = norm_type,
                     depthwise  = depthwise)
        )
        ## P5/32
        self.layer_5 = nn.Sequential(
            Conv(self.feat_dims[3], self.feat_dims[4], k=3, p=1, s=2, act_type=act_type, norm_type=norm_type),
            RTCBlock(in_dim     = self.feat_dims[4],
                     out_dim    = self.feat_dims[4],
                     num_blocks = round(3*depth),
                     shortcut   = True,
                     act_type   = act_type,
                     norm_type  = norm_type,
                     depthwise  = depthwise)
        )

    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 = [c3, c4, c5]

        return outputs


# ---------------------------- Functions ----------------------------
## Build Backbone network
def build_backbone(cfg, pretrained=False): 
    # build backbone model
    backbone = RTCBackbone(width=cfg['width'],
                           depth=cfg['depth'],
                           ratio=cfg['ratio'],
                           act_type=cfg['bk_act'],
                           norm_type=cfg['bk_norm'],
                           depthwise=cfg['bk_depthwise']
                           )
    feat_dims = backbone.feat_dims[-3:]

    # Model name
    width, depth, ratio = cfg['width'], cfg['depth'], cfg['ratio']
    model_name = "{}" if not cfg['bk_pretrained_mae'] else "mae_{}"
    if  width == 0.25   and depth == 0.34 and ratio == 2.0:
        model_name = model_name.format("rtcnet_n")
    elif width == 0.50  and depth == 0.34 and ratio == 2.0:
        model_name = model_name.format("rtcnet_s")
    elif width == 0.75  and depth == 0.67 and ratio == 1.5:
        model_name = model_name.format("rtcnet_m")
    elif width == 1.0   and depth == 1.0  and ratio == 1.0:
        model_name = model_name.format("rtcnet_l")
    elif width == 1.25  and depth == 1.34  and ratio == 1.0:
        model_name = model_name.format("rtcnet_x")
    else:
        raise NotImplementedError("No such model size : width={}, depth={}, ratio={}. ".format(width, depth, ratio))

    # Load pretrained weight
    if pretrained:
        backbone = load_pretrained_weight(backbone, model_name)
        
    return backbone, feat_dims

## Load pretrained weight
def load_pretrained_weight(model, model_name):
    # Load pretrained weight
    url = model_urls[model_name]
    if url is not None:
        print('Loading pretrained weight ...')
        checkpoint = torch.hub.load_state_dict_from_url(
            url=url, map_location="cpu", check_hash=True)
        # checkpoint state dict
        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:
                    checkpoint_state_dict.pop(k)
            else:
                checkpoint_state_dict.pop(k)
                print('Unused key: ', k)
        # load the weight
        model.load_state_dict(checkpoint_state_dict)
    else:
        print('No backbone pretrained for {}.'.format(model_name))

    return model


if __name__ == '__main__':
    import time
    from thop import profile
    cfg = {
        'bk_pretrained': True,
        'bk_pretrained_mae': False,
        'bk_act': 'silu',
        'bk_norm': 'BN',
        'bk_depthwise': False,
        'width': 0.25,
        'depth': 0.34,
        'ratio': 2.0,
    }
    model, feats = build_backbone(cfg, pretrained=cfg['bk_pretrained'])
    x = torch.randn(1, 3, 640, 640)
    t0 = time.time()
    outputs = model(x)
    t1 = time.time()
    print('Time: ', t1 - t0)
    for out in outputs:
        print(out.shape)

    x = torch.randn(1, 3, 640, 640)
    print('==============================')
    flops, params = profile(model, inputs=(x, ), verbose=False)
    print('==============================')
    print('GFLOPs : {:.2f}'.format(flops / 1e9 * 2))
    print('Params : {:.2f} M'.format(params / 1e6))