update

yolo/models/gelan/gelan_backbone.py

@@ -41,7 +41,7 @@ class GElanBackbone(nn.Module):
         # P3/8
         self.layer_3 = nn.Sequential(
             ADown(cfg.backbone_feats["c2"][2], cfg.backbone_feats["c3"][0],
-                  act_type=cfg.bk_act, norm_type=cfg.bk_norm, depthwise=cfg.bk_depthwise, use_pooling=cfg.bk_down_pooling),
+                  act_type=cfg.bk_act, norm_type=cfg.bk_norm, depthwise=cfg.bk_depthwise),
             RepGElanLayer(in_dim     = cfg.backbone_feats["c3"][0],
                           inter_dims = cfg.backbone_feats["c3"][1],
                           out_dim    = cfg.backbone_feats["c3"][2],
@@ -54,7 +54,7 @@ class GElanBackbone(nn.Module):
         # P4/16
         self.layer_4 = nn.Sequential(
             ADown(cfg.backbone_feats["c3"][2], cfg.backbone_feats["c4"][0],
-                  act_type=cfg.bk_act, norm_type=cfg.bk_norm, depthwise=cfg.bk_depthwise, use_pooling=cfg.bk_down_pooling),
+                  act_type=cfg.bk_act, norm_type=cfg.bk_norm, depthwise=cfg.bk_depthwise),
             RepGElanLayer(in_dim     = cfg.backbone_feats["c4"][0],
                           inter_dims = cfg.backbone_feats["c4"][1],
                           out_dim    = cfg.backbone_feats["c4"][2],
@@ -67,7 +67,7 @@ class GElanBackbone(nn.Module):
         # P5/32
         self.layer_5 = nn.Sequential(
             ADown(cfg.backbone_feats["c4"][2], cfg.backbone_feats["c5"][0],
-                  act_type=cfg.bk_act, norm_type=cfg.bk_norm, depthwise=cfg.bk_depthwise, use_pooling=cfg.bk_down_pooling),
+                  act_type=cfg.bk_act, norm_type=cfg.bk_norm, depthwise=cfg.bk_depthwise),
             RepGElanLayer(in_dim     = cfg.backbone_feats["c5"][0],
                           inter_dims = cfg.backbone_feats["c5"][1],
                           out_dim    = cfg.backbone_feats["c5"][2],

yolo/models/gelan/gelan_basic.py

@@ -72,30 +72,22 @@ class BasicConv(nn.Module):
 
 # --------------------- GELAN modules (from yolov9) ---------------------
 class ADown(nn.Module):
-    def __init__(self, in_dim, out_dim, act_type="silu", norm_type="BN", depthwise=False, use_pooling=True):
+    def __init__(self, in_dim, out_dim, act_type="silu", norm_type="BN", depthwise=False):
         super().__init__()
         inter_dim = out_dim // 2
-        self.use_pooling = use_pooling
-        if use_pooling:
-            self.conv_layer_1 = BasicConv(in_dim // 2, inter_dim,
-                                        kernel_size=3, padding=1, stride=2,
-                                        act_type=act_type, norm_type=norm_type, depthwise=depthwise)
-            self.conv_layer_2 = BasicConv(in_dim // 2, inter_dim, kernel_size=1,
-                                        act_type=act_type, norm_type=norm_type, depthwise=depthwise)
-        else:
-            self.conv_layer = BasicConv(in_dim, out_dim, kernel_size=3, padding=1, stride=2,
-                                        act_type=act_type, norm_type=norm_type, depthwise=depthwise)
+        self.conv_layer_1 = BasicConv(in_dim // 2, inter_dim,
+                                    kernel_size=3, padding=1, stride=2,
+                                    act_type=act_type, norm_type=norm_type, depthwise=depthwise)
+        self.conv_layer_2 = BasicConv(in_dim // 2, inter_dim, kernel_size=1,
+                                    act_type=act_type, norm_type=norm_type, depthwise=depthwise)
     def forward(self, x):
-        if self.use_pooling:
-            x = torch.nn.functional.avg_pool2d(x, 2, 1, 0, False, True)
-            x1,x2 = x.chunk(2, 1)
-            x1 = self.conv_layer_1(x1)
-            x2 = torch.nn.functional.max_pool2d(x2, 3, 2, 1)
-            x2 = self.conv_layer_2(x2)
+        x = torch.nn.functional.avg_pool2d(x, 2, 1, 0, False, True)
+        x1,x2 = x.chunk(2, 1)
+        x1 = self.conv_layer_1(x1)
+        x2 = torch.nn.functional.max_pool2d(x2, 3, 2, 1)
+        x2 = self.conv_layer_2(x2)
 
-            return torch.cat((x1, x2), 1)
-        else:
-            return self.conv_layer(x)
+        return torch.cat((x1, x2), 1)
 
 class RepConvN(nn.Module):
     """RepConv is a basic rep-style block, including training and deploy status

yolo/models/yolov7_af/yolov7_af_backbone.py

@@ -8,12 +8,12 @@ except:
 
 # IN1K pretrained weight
 pretrained_urls = {
-    't': None,
+    't': "https://github.com/yjh0410/ICLab/releases/download/in1k_pretrained/elannet_t_in1k_6.6.pth",
     'l': None,
     'x': None,
 }
 
-# ELANNet
+# ELANNet-Tiny
 class Yolov7TBackbone(nn.Module):
     def __init__(self, cfg):
         super(Yolov7TBackbone, self).__init__()
@@ -33,9 +33,14 @@ class Yolov7TBackbone(nn.Module):
         self.layer_4 = self.make_block(self.feat_dims[2], self.feat_dims[3], expansion=0.5)
         self.layer_5 = self.make_block(self.feat_dims[3], self.feat_dims[4], expansion=0.5)
 
+        # 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():
@@ -44,6 +49,31 @@ class Yolov7TBackbone(nn.Module):
                 # 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)
@@ -70,7 +100,7 @@ class Yolov7TBackbone(nn.Module):
 
         return outputs
 
-
+# ELANNet-Large
 class Yolov7LBackbone(nn.Module):
     def __init__(self, cfg):
         super(Yolov7LBackbone, self).__init__()
@@ -182,9 +212,9 @@ if __name__ == '__main__':
             self.bk_act = 'silu'
             self.bk_norm = 'BN'
             self.bk_depthwise = False
+            self.use_pretrained = True
             self.width = 0.5
-            self.depth = 0.34
-            self.scale = "l"
+            self.scale = "t"
 
     cfg = BaseConfig()
     model = Yolov7TBackbone(cfg)