use in1k pretrained for coco

yolo/config/yolov3_config.py

@@ -29,6 +29,7 @@ class Yolov3BaseConfig(object):
         self.bk_act   = 'silu'
         self.bk_norm  = 'BN'
         self.bk_depthwise = False
+        self.use_pretrained = True
         ## Neck
         self.neck_act       = 'silu'
         self.neck_norm      = 'BN'

yolo/config/yolov5_af_config.py

@@ -29,6 +29,7 @@ class Yolov5AFBaseConfig(object):
         self.bk_act   = 'silu'
         self.bk_norm  = 'BN'
         self.bk_depthwise = False
+        self.use_pretrained = True
         ## Neck
         self.neck_act       = 'silu'
         self.neck_norm      = 'BN'

yolo/config/yolov5_config.py

@@ -29,6 +29,7 @@ class Yolov5BaseConfig(object):
         self.bk_act   = 'silu'
         self.bk_norm  = 'BN'
         self.bk_depthwise = False
+        self.use_pretrained = True
         ## Neck
         self.neck_act       = 'silu'
         self.neck_norm      = 'BN'

yolo/config/yolov7_af_config.py

@@ -22,7 +22,7 @@ class Yolov7AFBaseConfig(object):
         self.bk_act   = 'silu'
         self.bk_norm  = 'BN'
         self.bk_depthwise = False
-        self.use_pretrained = False
+        self.use_pretrained = True
         ## Neck
         self.neck_act       = 'silu'
         self.neck_norm      = 'BN'

yolo/config/yolov8_config.py

@@ -31,7 +31,7 @@ class Yolov8BaseConfig(object):
         self.bk_act   = 'silu'
         self.bk_norm  = 'BN'
         self.bk_depthwise = False
-        self.use_pretrained = False
+        self.use_pretrained = True
         ## Neck
         self.neck_act       = 'silu'
         self.neck_norm      = 'BN'

yolo/models/yolov3/yolov3_backbone.py

@@ -6,6 +6,14 @@ try:
 except:
     from  yolov3_basic import BasicConv, ResBlock
 
+# IN1K pretrained weight
+pretrained_urls = {
+    'n': None,
+    's': None,
+    'm': None,
+    'l': None,
+    'x': None,
+}
 
 # --------------------- Yolov3's Backbone -----------------------
 ## Modified DarkNet
@@ -84,7 +92,11 @@ class Yolov3Backbone(nn.Module):
 
         # 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():
@@ -93,6 +105,31 @@ class Yolov3Backbone(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 forward(self, x):
         c1 = self.layer_1(x)
         c2 = self.layer_2(c1)

yolo/models/yolov5/yolov5_backbone.py

@@ -6,6 +6,14 @@ try:
 except:
     from  yolov5_basic import BasicConv, CSPBlock
 
+# IN1K pretrained weight
+pretrained_urls = {
+    'n': None,
+    's': None,
+    'm': None,
+    'l': None,
+    'x': None,
+}
 
 # --------------------- Yolov3's Backbone -----------------------
 ## Modified DarkNet
@@ -84,6 +92,10 @@ class Yolov5Backbone(nn.Module):
 
         # Initialize all layers
         self.init_weights()
+
+        # Load imagenet pretrained weight
+        if cfg.use_pretrained:
+            self.load_pretrained()
         
     def init_weights(self):
         """Initialize the parameters."""
@@ -93,6 +105,31 @@ class Yolov5Backbone(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 forward(self, x):
         c1 = self.layer_1(x)
         c2 = self.layer_2(c1)

yolo/models/yolov5_af/yolov5_af_backbone.py

@@ -6,6 +6,14 @@ try:
 except:
     from  yolov5_af_basic import BasicConv, CSPBlock
 
+# IN1K pretrained weight
+pretrained_urls = {
+    'n': None,
+    's': None,
+    'm': None,
+    'l': None,
+    'x': None,
+}
 
 # --------------------- Yolov3's Backbone -----------------------
 ## Modified DarkNet
@@ -85,6 +93,10 @@ class Yolov5Backbone(nn.Module):
         # 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():
@@ -93,6 +105,31 @@ class Yolov5Backbone(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 forward(self, x):
         c1 = self.layer_1(x)
         c2 = self.layer_2(c1)

yolo/models/yolov7_af/yolov7_af_backbone.py

@@ -6,6 +6,12 @@ try:
 except:
     from  yolov7_af_basic import BasicConv, MDown, ELANLayer
 
+# IN1K pretrained weight
+pretrained_urls = {
+    't': None,
+    'l': None,
+    'x': None,
+}
 
 # ELANNet
 class Yolov7TBackbone(nn.Module):
@@ -87,6 +93,10 @@ class Yolov7LBackbone(nn.Module):
         # 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():
@@ -95,6 +105,31 @@ class Yolov7LBackbone(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 = nn.Sequential(
             BasicConv(in_dim, out_dim//2, kernel_size=3, padding=1, stride=1,
@@ -147,12 +182,12 @@ if __name__ == '__main__':
             self.bk_act = 'silu'
             self.bk_norm = 'BN'
             self.bk_depthwise = False
-            self.width = 1.0
+            self.width = 0.5
             self.depth = 0.34
             self.scale = "l"
 
     cfg = BaseConfig()
-    model = Yolov7LBackbone(cfg)
+    model = Yolov7TBackbone(cfg)
     x = torch.randn(1, 3, 640, 640)
     t0 = time.time()
     outputs = model(x)

yolo/models/yolov8/yolov8_backbone.py

@@ -6,6 +6,14 @@ try:
 except:
     from  yolov8_basic import BasicConv, ELANLayer
 
+# IN1K pretrained weight
+pretrained_urls = {
+    'n': "https://github.com/yjh0410/ICLab/releases/download/in1k_pretrained/rtcnet_n_in1k_62.1.pth",
+    's': "https://github.com/yjh0410/ICLab/releases/download/in1k_pretrained/rtcnet_s_in1k_71.3.pth",
+    'm': None,
+    'l': None,
+    'x': None,
+}
 
 # ---------------------------- Basic functions ----------------------------
 class Yolov8Backbone(nn.Module):
@@ -84,6 +92,10 @@ class Yolov8Backbone(nn.Module):
         # 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():
@@ -92,6 +104,31 @@ class Yolov8Backbone(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 forward(self, x):
         c1 = self.layer_1(x)
         c2 = self.layer_2(c1)
@@ -120,10 +157,11 @@ if __name__ == '__main__':
             self.bk_act = 'silu'
             self.bk_norm = 'BN'
             self.bk_depthwise = False
-            self.width = 1.0
-            self.depth = 1.0
-            self.ratio = 1.0
-            self.scale = "n"
+            self.use_pretrained = True
+            self.width = 0.50
+            self.depth = 0.34
+            self.ratio = 2.0
+            self.scale = "s"
 
     cfg = BaseConfig()
     model = build_backbone(cfg)