modify YOLOv7-AF

yolo/config/yolov7_af_config.py

@@ -124,6 +124,7 @@ class Yolov7AFTConfig(Yolov7AFBaseConfig):
         # ---------------- Model config ----------------
         self.width = 0.50
         self.scale = "t"
+        self.use_pretrained = True
         self.fpn_expansions = [0.5, 0.5]
         self.fpn_block_bw = 2
         self.fpn_block_dw = 1

yolo/config/yolov8_config.py

@@ -88,7 +88,7 @@ class Yolov8BaseConfig(object):
         # ---------------- Lr Scheduler config ----------------
         self.warmup_epoch = 3
         self.lr_scheduler = "cosine"
-        self.max_epoch    = 500
+        self.max_epoch    = 300
         self.eval_epoch   = 10
         self.no_aug_epoch = 20
 

yolo/models/yolov7_af/yolov7_af.py

@@ -4,7 +4,7 @@ import torch.nn as nn
 
 # --------------- Model components ---------------
 from .yolov7_af_backbone import Yolov7TBackbone, Yolov7LBackbone
-from .yolov7_af_neck     import SPPF
+from .yolov7_af_neck     import SPPFBlockCSP
 from .yolov7_af_pafpn    import Yolov7PaFPN
 from .yolov7_af_head     import Yolov7DetHead
 from .yolov7_af_pred     import Yolov7AFDetPredLayer
@@ -35,7 +35,7 @@ class Yolov7AF(nn.Module):
         self.backbone = Yolov7TBackbone(cfg) if cfg.scale == "t" else Yolov7LBackbone(cfg)
         self.pyramid_feat_dims = self.backbone.feat_dims[-3:]
         ## Neck: SPP
-        self.neck     = SPPF(cfg, self.pyramid_feat_dims[-1], self.pyramid_feat_dims[-1]//2)
+        self.neck     = SPPFBlockCSP(cfg, self.pyramid_feat_dims[-1], self.pyramid_feat_dims[-1]//2)
         self.pyramid_feat_dims[-1] = self.neck.out_dim
         ## Neck: FPN
         self.fpn      = Yolov7PaFPN(cfg, self.pyramid_feat_dims)

yolo/models/yolov7_af/yolov7_af_backbone.py

@@ -28,10 +28,10 @@ class Yolov7TBackbone(nn.Module):
 
         # ---------------- 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)
-        self.layer_3 = self.make_block(self.feat_dims[1], self.feat_dims[2], expansion=0.5)
-        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)
+        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
@@ -80,13 +80,22 @@ class Yolov7TBackbone(nn.Module):
         
         return stem
 
-    def make_block(self, in_dim, out_dim, expansion=0.5):
-        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),
-                    )
+    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
     
@@ -212,7 +221,7 @@ if __name__ == '__main__':
             self.bk_act = 'silu'
             self.bk_norm = 'BN'
             self.bk_depthwise = False
-            self.use_pretrained = True
+            self.use_pretrained = False
             self.width = 0.5
             self.scale = "t"
 

yolo/models/yolov7_af/yolov7_af_pafpn.py

@@ -15,7 +15,6 @@ class Yolov7PaFPN(nn.Module):
         super(Yolov7PaFPN, self).__init__()
         # ----------------------------- Basic parameters -----------------------------
         self.in_dims = in_dims
-        self.out_dims = [round(256*cfg.width), round(512*cfg.width), round(1024*cfg.width)]
         c3, c4, c5 = in_dims
 
         # ----------------------------- Yolov7's Top-down FPN -----------------------------
@@ -85,6 +84,14 @@ class Yolov7PaFPN(nn.Module):
                                      kernel_size=3, padding=1, stride=1,
                                      act_type=cfg.fpn_act, norm_type=cfg.fpn_norm, depthwise=cfg.fpn_depthwise)
 
+        # ---------------------- Yolox's output projection ----------------------
+        self.out_layers = nn.ModuleList([
+            BasicConv(in_dim, round(cfg.head_dim*cfg.width), kernel_size=1,
+                      act_type=cfg.fpn_act, norm_type=cfg.fpn_norm)
+                      for in_dim in [round(256*cfg.width), round(512*cfg.width), round(1024*cfg.width)]
+                      ])
+        self.out_dims = [round(cfg.head_dim*cfg.width)] * 3
+
         # Initialize all layers
         self.init_weights()
 
@@ -125,7 +132,12 @@ class Yolov7PaFPN(nn.Module):
 
         out_feats = [self.head_conv_1(p3), self.head_conv_2(p4), self.head_conv_3(p5)]
             
-        return out_feats
+        # output proj layers
+        out_feats_proj = []
+        for feat, layer in zip(out_feats, self.out_layers):
+            out_feats_proj.append(layer(feat))
+            
+        return out_feats_proj
 
 
 if __name__=='__main__':