YOLO-Tutorial-v2/yolo/models/yolov7_af/yolov7_af_backbone.py

import torch
import torch.nn as nn

try:
    from .yolov7_af_basic import BasicConv, MDown, ELANLayer
except:
    from  yolov7_af_basic import BasicConv, MDown, ELANLayer

# IN1K pretrained weight
pretrained_urls = {
    't': "https://github.com/yjh0410/ICLab/releases/download/in1k_pretrained/elannet_t_in1k_63.2.pth",
    'l': None,
    'x': None,
}

# ELANNet-Tiny
class Yolov7TBackbone(nn.Module):
    def __init__(self, cfg):
        super(Yolov7TBackbone, self).__init__()
        # ---------------- Basic parameters ----------------
        self.model_scale = cfg.scale
        self.bk_act = cfg.bk_act
        self.bk_norm = cfg.bk_norm
        self.bk_depthwise = cfg.bk_depthwise
        self.elan_depth = 1
        self.feat_dims = [round(64  * cfg.width), round(128  * cfg.width),
                          round(256  * cfg.width), round(512 * cfg.width), round(1024 * cfg.width)]

        # ---------------- Model parameters ----------------
        self.layer_1 = self.make_stem(3, self.feat_dims[0])
        self.layer_2 = self.make_block(self.feat_dims[0], self.feat_dims[1], expansion=0.5, downsample="conv")
        self.layer_3 = self.make_block(self.feat_dims[1], self.feat_dims[2], expansion=0.5, downsample="maxpool")
        self.layer_4 = self.make_block(self.feat_dims[2], self.feat_dims[3], expansion=0.5, downsample="maxpool")
        self.layer_5 = self.make_block(self.feat_dims[3], self.feat_dims[4], expansion=0.5, downsample="maxpool")

        # Initialize all layers
        # Initialize all layers
        self.init_weights()
        
        # Load imagenet pretrained weight
        if cfg.use_pretrained:
            self.load_pretrained()
        
    def init_weights(self):
        """Initialize the parameters."""
        for m in self.modules():
            if isinstance(m, torch.nn.Conv2d):
                # In order to be consistent with the source code,
                # reset the Conv2d initialization parameters
                m.reset_parameters()

    def load_pretrained(self):
        url = pretrained_urls[self.model_scale]
        if url is not None:
            print('Loading backbone pretrained weight from : {}'.format(url))
            # checkpoint state dict
            checkpoint = torch.hub.load_state_dict_from_url(
                url=url, map_location="cpu", check_hash=True)
            checkpoint_state_dict = checkpoint.pop("model")
            # model state dict
            model_state_dict = self.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
            self.load_state_dict(checkpoint_state_dict)
        else:
            print('No pretrained weight for model scale: {}.'.format(self.model_scale))

    def make_stem(self, in_dim, out_dim):
        stem = BasicConv(in_dim, out_dim, kernel_size=6, padding=2, stride=2,
                         act_type=self.bk_act, norm_type=self.bk_norm, depthwise=self.bk_depthwise)
        
        return stem

    def make_block(self, in_dim, out_dim, expansion=0.5, downsample="maxpool"):
        if downsample == "maxpool":
            block = nn.Sequential(
                nn.MaxPool2d((2, 2), stride=2),             
                ELANLayer(in_dim, out_dim, expansion=expansion, num_blocks=self.elan_depth,
                          act_type=self.bk_act, norm_type=self.bk_norm, depthwise=self.bk_depthwise),
                          )
        elif downsample == "conv":
            block = nn.Sequential(
                BasicConv(in_dim, out_dim, kernel_size=3, padding=1, stride=2,
                          act_type=self.bk_act, norm_type=self.bk_norm, depthwise=self.bk_depthwise),             
                ELANLayer(out_dim, out_dim, expansion=expansion, num_blocks=self.elan_depth,
                          act_type=self.bk_act, norm_type=self.bk_norm, depthwise=self.bk_depthwise),
                          )
        else:
            raise NotImplementedError("Unknown downsample type: {}".format(downsample))

        return block
    
    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

# ELANNet-Large
class Yolov7LBackbone(nn.Module):
    def __init__(self, cfg):
        super(Yolov7LBackbone, self).__init__()
        # ---------------- Basic parameters ----------------
        self.model_scale = cfg.scale
        self.bk_act = cfg.bk_act
        self.bk_norm = cfg.bk_norm
        self.bk_depthwise = cfg.bk_depthwise
        self.elan_depth = 2
        self.feat_dims = [round(64  * cfg.width), round(128  * cfg.width), round(256  * cfg.width),
                          round(512  * cfg.width), round(1024 * cfg.width), round(1024 * cfg.width)]

        # ---------------- Model parameters ----------------
        self.layer_1 = self.make_stem(3, self.feat_dims[0])
        self.layer_2 = self.make_block(self.feat_dims[0], self.feat_dims[1], self.feat_dims[2], expansion=0.5, conv_downsample=True)
        self.layer_3 = self.make_block(self.feat_dims[2], self.feat_dims[2], self.feat_dims[3], expansion=0.5)
        self.layer_4 = self.make_block(self.feat_dims[3], self.feat_dims[3], self.feat_dims[4], expansion=0.5)
        self.layer_5 = self.make_block(self.feat_dims[4], self.feat_dims[4], self.feat_dims[5], expansion=0.25)

        # Initialize all layers
        self.init_weights()
        
        # Load imagenet pretrained weight
        if cfg.use_pretrained:
            self.load_pretrained()
        
    def init_weights(self):
        """Initialize the parameters."""
        for m in self.modules():
            if isinstance(m, torch.nn.Conv2d):
                # In order to be consistent with the source code,
                # reset the Conv2d initialization parameters
                m.reset_parameters()

    def load_pretrained(self):
        url = pretrained_urls[self.model_scale]
        if url is not None:
            print('Loading backbone pretrained weight from : {}'.format(url))
            # checkpoint state dict
            checkpoint = torch.hub.load_state_dict_from_url(
                url=url, map_location="cpu", check_hash=True)
            checkpoint_state_dict = checkpoint.pop("model")
            # model state dict
            model_state_dict = self.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
            self.load_state_dict(checkpoint_state_dict)
        else:
            print('No pretrained weight for model scale: {}.'.format(self.model_scale))

    def make_stem(self, in_dim, out_dim):
        stem = nn.Sequential(
            BasicConv(in_dim, out_dim//2, kernel_size=3, padding=1, stride=1,
                        act_type=self.bk_act, norm_type=self.bk_norm, depthwise=self.bk_depthwise),
            BasicConv(out_dim//2, out_dim, kernel_size=3, padding=1, stride=2,
                        act_type=self.bk_act, norm_type=self.bk_norm, depthwise=self.bk_depthwise),
            BasicConv(out_dim, out_dim, kernel_size=3, padding=1, stride=1,
                        act_type=self.bk_act, norm_type=self.bk_norm, depthwise=self.bk_depthwise)

        )

        return stem

    def make_block(self, in_dim, out_dim_1, out_dim_2, expansion=0.5, conv_downsample=False):
        if conv_downsample:
            block = nn.Sequential(
                BasicConv(in_dim, out_dim_1, kernel_size=3, padding=1, stride=2,
                         act_type=self.bk_act, norm_type=self.bk_norm, depthwise=self.bk_depthwise),             
                ELANLayer(out_dim_1, out_dim_2,
                        expansion=expansion, num_blocks=self.elan_depth,
                        act_type=self.bk_act, norm_type=self.bk_norm, depthwise=self.bk_depthwise),
        )
        else:
            block = nn.Sequential(
                MDown(in_dim, out_dim_1,
                    act_type=self.bk_act, norm_type=self.bk_norm, depthwise=self.bk_depthwise),             
                ELANLayer(out_dim_1, out_dim_2,
                        expansion=expansion, num_blocks=self.elan_depth,
                        act_type=self.bk_act, norm_type=self.bk_norm, depthwise=self.bk_depthwise),
            )
        
        return block
    
    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


if __name__ == '__main__':
    import time
    from thop import profile
    class BaseConfig(object):
        def __init__(self) -> None:
            self.bk_act = 'silu'
            self.bk_norm = 'BN'
            self.bk_depthwise = False
            self.use_pretrained = False
            self.width = 0.5
            self.scale = "t"

    cfg = BaseConfig()
    model = Yolov7TBackbone(cfg)
    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))