optimize codes style

config/yolov1_config.py

@@ -36,7 +36,7 @@ class Yolov1BaseConfig(object):
         self.val_conf_thresh = 0.001
         self.val_nms_thresh  = 0.7
         self.test_topk = 100
-        self.test_conf_thresh = 0.2
+        self.test_conf_thresh = 0.3
         self.test_nms_thresh  = 0.5
 
         # ---------------- Assignment config ----------------

config/yolov2_config.py

@@ -37,7 +37,7 @@ class Yolov2BaseConfig(object):
         self.val_conf_thresh = 0.001
         self.val_nms_thresh  = 0.7
         self.test_topk = 100
-        self.test_conf_thresh = 0.2
+        self.test_conf_thresh = 0.3
         self.test_nms_thresh  = 0.5
 
         # ---------------- Assignment config ----------------

models/yolov1/yolov1.py

@@ -32,7 +32,7 @@ class Yolov1(nn.Module):
         self.backbone = Yolov1Backbone(cfg)
         self.neck     = SPPF(cfg, self.backbone.feat_dim, cfg.head_dim)
         self.head     = Yolov1DetHead(cfg, self.neck.out_dim)
-        self.pred     = Yolov1DetPredLayer(cfg, self.num_classes)
+        self.pred     = Yolov1DetPredLayer(cfg)
 
     def post_process(self, obj_preds, cls_preds, box_preds):
         """

models/yolov1/yolov1_backbone.py

@@ -44,7 +44,6 @@ if __name__=='__main__':
     print('Time: ', t1 - t0)
     print(output.shape)
 
-    print('==============================')
     flops, params = profile(model, inputs=(x, ), verbose=False)
     print('==============================')
     print('GFLOPs : {:.2f}'.format(flops / 1e9 * 2))

models/yolov1/yolov1_head.py

@@ -114,7 +114,6 @@ if __name__=='__main__':
     print('Time: ', t1 - t0)
     print(cls_feat.shape, reg_feat.shape)
 
-    print('==============================')
     flops, params = profile(head, inputs=(x, ), verbose=False)
     print('==============================')
     print('GFLOPs : {:.2f}'.format(flops / 1e9 * 2))

models/yolov1/yolov1_neck.py

@@ -1,7 +1,10 @@
 import torch
 import torch.nn as nn
 
-from .yolov1_basic import BasicConv
+try:
+    from .yolov1_basic import BasicConv
+except:
+    from  yolov1_basic import BasicConv
 
 
 # Spatial Pyramid Pooling - Fast (SPPF) layer for YOLOv5 by Glenn Jocher
@@ -31,3 +34,41 @@ class SPPF(nn.Module):
         y2 = self.m(y1)
 
         return self.cv2(torch.cat((x, y1, y2, self.m(y2)), 1))
+
+
+if __name__=='__main__':
+    import time
+    from thop import profile
+    # Model config
+    
+    # YOLOv8-Base config
+    class Yolov1BaseConfig(object):
+        def __init__(self) -> None:
+            # ---------------- Model config ----------------
+            self.out_stride = 32
+            self.max_stride = 32
+            ## Neck
+            self.neck_act       = 'lrelu'
+            self.neck_norm      = 'BN'
+            self.neck_depthwise = False
+            self.neck_expand_ratio = 0.5
+            self.spp_pooling_size  = 5
+
+    cfg = Yolov1BaseConfig()
+    # Build a head
+    in_dim  = 512
+    out_dim = 512
+    neck = SPPF(cfg, 512, 512)
+
+    # Inference
+    x = torch.randn(1, in_dim, 20, 20)
+    t0 = time.time()
+    output = neck(x)
+    t1 = time.time()
+    print('Time: ', t1 - t0)
+    print('Neck output: ', output.shape)
+
+    flops, params = profile(neck, inputs=(x, ), verbose=False)
+    print('==============================')
+    print('GFLOPs : {:.2f}'.format(flops / 1e9 * 2))
+    print('Params : {:.2f} M'.format(params / 1e6))

models/yolov1/yolov1_pred.py

@@ -5,16 +5,17 @@ import torch.nn as nn
 # -------------------- Detection Pred Layer --------------------
 ## Single-level pred layer
 class Yolov1DetPredLayer(nn.Module):
-    def __init__(self, cfg, num_classes):
+    def __init__(self, cfg):
         super().__init__()
         # --------- Basic Parameters ----------
         self.stride  = cfg.out_stride
         self.cls_dim = cfg.head_dim
         self.reg_dim = cfg.head_dim
+        self.num_classes = cfg.num_classes
 
         # --------- Network Parameters ----------
         self.obj_pred = nn.Conv2d(self.cls_dim, 1, kernel_size=1)
-        self.cls_pred = nn.Conv2d(self.cls_dim, num_classes, kernel_size=1)
+        self.cls_pred = nn.Conv2d(self.cls_dim, self.num_classes, kernel_size=1)
         self.reg_pred = nn.Conv2d(self.reg_dim, 4, kernel_size=1)                
 
         self.init_bias()
@@ -89,3 +90,42 @@ class Yolov1DetPredLayer(nn.Module):
                    }
 
         return outputs
+
+
+if __name__=='__main__':
+    import time
+    from thop import profile
+    # Model config
+    
+    # YOLOv8-Base config
+    class Yolov1BaseConfig(object):
+        def __init__(self) -> None:
+            # ---------------- Model config ----------------
+            self.out_stride = 32
+            self.max_stride = 32
+            ## Head
+            self.head_dim  = 512
+
+    cfg = Yolov1BaseConfig()
+    cfg.num_classes = 20
+    # Build a pred layer
+    pred = Yolov1DetPredLayer(cfg)
+
+    # Inference
+    cls_feat = torch.randn(1, cfg.head_dim, 20, 20)
+    reg_feat = torch.randn(1, cfg.head_dim, 20, 20)
+    t0 = time.time()
+    output = pred(cls_feat, reg_feat)
+    t1 = time.time()
+    print('Time: ', t1 - t0)
+    print('====== Pred output ======= ')
+    for k in output:
+        if isinstance(output[k], torch.Tensor):
+            print("-{}: ".format(k), output[k].shape)
+        else:
+            print("-{}: ".format(k), output[k])
+
+    flops, params = profile(pred, inputs=(cls_feat, reg_feat, ), verbose=False)
+    print('==============================')
+    print('GFLOPs : {:.2f}'.format(flops / 1e9 * 2))
+    print('Params : {:.2f} M'.format(params / 1e6))

models/yolov2/yolov2.py

@@ -32,7 +32,7 @@ class Yolov2(nn.Module):
         self.backbone = Yolov2Backbone(cfg)
         self.neck     = SPPF(cfg, self.backbone.feat_dim, cfg.head_dim)
         self.head     = Yolov2DetHead(cfg, self.neck.out_dim)
-        self.pred     = Yolov2DetPredLayer(cfg, self.num_classes)
+        self.pred     = Yolov2DetPredLayer(cfg)
 
     def post_process(self, obj_preds, cls_preds, box_preds):
         """

models/yolov2/yolov2_head.py

@@ -88,7 +88,7 @@ if __name__=='__main__':
     # Model config
     
     # YOLOv8-Base config
-    class Yolov1BaseConfig(object):
+    class Yolov2BaseConfig(object):
         def __init__(self) -> None:
             # ---------------- Model config ----------------
             self.out_stride = 32
@@ -101,7 +101,7 @@ if __name__=='__main__':
             self.num_cls_head   = 2
             self.num_reg_head   = 2
 
-    cfg = Yolov1BaseConfig()
+    cfg = Yolov2BaseConfig()
     # Build a head
     head = Yolov2DetHead(cfg, 512)
 
@@ -118,4 +118,4 @@ if __name__=='__main__':
     flops, params = profile(head, inputs=(x, ), verbose=False)
     print('==============================')
     print('GFLOPs : {:.2f}'.format(flops / 1e9 * 2))
-    print('Params : {:.2f} M'.format(params / 1e6))    
+    print('Params : {:.2f} M'.format(params / 1e6))    

models/yolov2/yolov2_neck.py

@@ -1,7 +1,10 @@
 import torch
 import torch.nn as nn
 
-from .yolov2_basic import BasicConv
+try:
+    from .yolov2_basic import BasicConv
+except:
+    from  yolov2_basic import BasicConv
 
 
 # Spatial Pyramid Pooling - Fast (SPPF) layer for YOLOv5 by Glenn Jocher
@@ -31,3 +34,41 @@ class SPPF(nn.Module):
         y2 = self.m(y1)
 
         return self.cv2(torch.cat((x, y1, y2, self.m(y2)), 1))
+
+
+if __name__=='__main__':
+    import time
+    from thop import profile
+    # Model config
+    
+    # YOLOv2-Base config
+    class Yolov2BaseConfig(object):
+        def __init__(self) -> None:
+            # ---------------- Model config ----------------
+            self.out_stride = 32
+            self.max_stride = 32
+            ## Neck
+            self.neck_act       = 'lrelu'
+            self.neck_norm      = 'BN'
+            self.neck_depthwise = False
+            self.neck_expand_ratio = 0.5
+            self.spp_pooling_size  = 5
+
+    cfg = Yolov2BaseConfig()
+    # Build a head
+    in_dim  = 512
+    out_dim = 512
+    neck = SPPF(cfg, in_dim, out_dim)
+
+    # Inference
+    x = torch.randn(1, in_dim, 20, 20)
+    t0 = time.time()
+    output = neck(x)
+    t1 = time.time()
+    print('Time: ', t1 - t0)
+    print('Neck output: ', output.shape)
+
+    flops, params = profile(neck, inputs=(x, ), verbose=False)
+    print('==============================')
+    print('GFLOPs : {:.2f}'.format(flops / 1e9 * 2))
+    print('Params : {:.2f} M'.format(params / 1e6))

models/yolov2/yolov2_pred.py

@@ -5,7 +5,7 @@ import torch.nn as nn
 # -------------------- Detection Pred Layer --------------------
 ## Single-level pred layer
 class Yolov2DetPredLayer(nn.Module):
-    def __init__(self, cfg, num_classes):
+    def __init__(self, cfg):
         super().__init__()
         # --------- Basic Parameters ----------
         self.stride  = cfg.out_stride
@@ -18,7 +18,7 @@ class Yolov2DetPredLayer(nn.Module):
 
         # --------- Network Parameters ----------
         self.obj_pred = nn.Conv2d(self.cls_dim, 1 * self.num_anchors, kernel_size=1)
-        self.cls_pred = nn.Conv2d(self.cls_dim, num_classes * self.num_anchors, kernel_size=1)
+        self.cls_pred = nn.Conv2d(self.cls_dim, self.num_classes * self.num_anchors, kernel_size=1)
         self.reg_pred = nn.Conv2d(self.reg_dim, 4 * self.num_anchors, kernel_size=1)                
 
         self.init_bias()
@@ -100,3 +100,43 @@ class Yolov2DetPredLayer(nn.Module):
                    }
 
         return outputs
+
+
+if __name__=='__main__':
+    import time
+    from thop import profile
+    # Model config
+    
+    # YOLOv8-Base config
+    class Yolov2BaseConfig(object):
+        def __init__(self) -> None:
+            # ---------------- Model config ----------------
+            self.out_stride = 32
+            self.max_stride = 32
+            ## Head
+            self.head_dim  = 512
+            self.anchor_sizes = [[17, 25], [55, 75], [92, 206], [202, 21], [289, 311]]
+
+    cfg = Yolov2BaseConfig()
+    cfg.num_classes = 20
+    # Build a pred layer
+    pred = Yolov2DetPredLayer(cfg)
+
+    # Inference
+    cls_feat = torch.randn(1, cfg.head_dim, 20, 20)
+    reg_feat = torch.randn(1, cfg.head_dim, 20, 20)
+    t0 = time.time()
+    output = pred(cls_feat, reg_feat)
+    t1 = time.time()
+    print('Time: ', t1 - t0)
+    print('====== Pred output ======= ')
+    for k in output:
+        if isinstance(output[k], torch.Tensor):
+            print("-{}: ".format(k), output[k].shape)
+        else:
+            print("-{}: ".format(k), output[k])
+
+    flops, params = profile(pred, inputs=(cls_feat, reg_feat, ), verbose=False)
+    print('==============================')
+    print('GFLOPs : {:.2f}'.format(flops / 1e9 * 2))
+    print('Params : {:.2f} M'.format(params / 1e6))

models/yolov3/yolov3_backbone.py

@@ -128,7 +128,6 @@ if __name__ == '__main__':
         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))

models/yolov3/yolov3_fpn.py

@@ -3,7 +3,10 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 
-from .yolov3_basic import BasicConv, ResBlock
+try:
+    from .yolov3_basic import BasicConv, ResBlock
+except:
+    from  yolov3_basic import BasicConv, ResBlock
 
 
 # Yolov3FPN
@@ -77,3 +80,47 @@ class Yolov3FPN(nn.Module):
             out_feats_proj.append(layer(feat))
 
         return out_feats_proj
+
+
+if __name__=='__main__':
+    import time
+    from thop import profile
+    # Model config
+    
+    # YOLOv2-Base config
+    class Yolov3BaseConfig(object):
+        def __init__(self) -> None:
+            # ---------------- Model config ----------------
+            self.width    = 0.50
+            self.depth    = 0.34
+            self.out_stride = [8, 16, 32]
+            self.max_stride = 32
+            self.num_levels = 3
+            ## FPN
+            self.fpn_act  = 'silu'
+            self.fpn_norm = 'BN'
+            self.fpn_depthwise = False
+            ## Head
+            self.head_dim = 256
+
+    cfg = Yolov3BaseConfig()
+    # Build a head
+    in_dims  = [128, 256, 512]
+    fpn = Yolov3FPN(cfg, in_dims)
+
+    # Inference
+    x = [torch.randn(1, in_dims[0], 80, 80),
+         torch.randn(1, in_dims[1], 40, 40),
+         torch.randn(1, in_dims[2], 20, 20)]
+    t0 = time.time()
+    output = fpn(x)
+    t1 = time.time()
+    print('Time: ', t1 - t0)
+    print('====== FPN output ====== ')
+    for level, feat in enumerate(output):
+        print("- Level-{} : ".format(level), feat.shape)
+
+    flops, params = profile(fpn, inputs=(x, ), verbose=False)
+    print('==============================')
+    print('GFLOPs : {:.2f}'.format(flops / 1e9 * 2))
+    print('Params : {:.2f} M'.format(params / 1e6))

models/yolov3/yolov3_head.py

@@ -137,7 +137,9 @@ if __name__=='__main__':
     class Yolov3BaseConfig(object):
         def __init__(self) -> None:
             # ---------------- Model config ----------------
-            self.out_stride = 32
+            self.width    = 0.50
+            self.depth    = 0.34
+            self.out_stride = [8, 16, 32]
             self.max_stride = 32
             self.num_levels = 3
             ## Head
@@ -161,10 +163,10 @@ if __name__=='__main__':
     cls_feats, reg_feats = head(pyramid_feats)
     t1 = time.time()
     print('Time: ', t1 - t0)
-    for cls_f, reg_f in zip(cls_feats, reg_feats):
-        print(cls_f.shape, reg_f.shape)
+    print("====== Yolov3 Head output ======")
+    for level, (cls_f, reg_f) in enumerate(zip(cls_feats, reg_feats)):
+        print("- Level-{} : ".format(level), cls_f.shape, reg_f.shape)
 
-    print('==============================')
     flops, params = profile(head, inputs=(pyramid_feats, ), verbose=False)
     print('==============================')
     print('GFLOPs : {:.2f}'.format(flops / 1e9 * 2))

models/yolov3/yolov3_neck.py

@@ -1,7 +1,10 @@
 import torch
 import torch.nn as nn
 
-from .yolov3_basic import BasicConv
+try:
+    from .yolov3_basic import BasicConv
+except:
+    from  yolov3_basic import BasicConv
 
 
 # Spatial Pyramid Pooling - Fast (SPPF) layer for YOLOv5 by Glenn Jocher
@@ -31,3 +34,41 @@ class SPPF(nn.Module):
         y2 = self.m(y1)
 
         return self.cv2(torch.cat((x, y1, y2, self.m(y2)), 1))
+
+
+if __name__=='__main__':
+    import time
+    from thop import profile
+    # Model config
+    
+    # YOLOv3-Base config
+    class Yolov3BaseConfig(object):
+        def __init__(self) -> None:
+            # ---------------- Model config ----------------
+            self.out_stride = 32
+            self.max_stride = 32
+            ## Neck
+            self.neck_act       = 'lrelu'
+            self.neck_norm      = 'BN'
+            self.neck_depthwise = False
+            self.neck_expand_ratio = 0.5
+            self.spp_pooling_size  = 5
+
+    cfg = Yolov3BaseConfig()
+    # Build a head
+    in_dim  = 512
+    out_dim = 512
+    neck = SPPF(cfg, in_dim, out_dim)
+
+    # Inference
+    x = torch.randn(1, in_dim, 20, 20)
+    t0 = time.time()
+    output = neck(x)
+    t1 = time.time()
+    print('Time: ', t1 - t0)
+    print('Neck output: ', output.shape)
+
+    flops, params = profile(neck, inputs=(x, ), verbose=False)
+    print('==============================')
+    print('GFLOPs : {:.2f}'.format(flops / 1e9 * 2))
+    print('Params : {:.2f} M'.format(params / 1e6))

models/yolov3/yolov3_pred.py

@@ -156,3 +156,59 @@ class Yolov3DetPredLayer(nn.Module):
                    }
 
         return outputs
+
+
+if __name__=='__main__':
+    import time
+    from thop import profile
+    # Model config
+    
+    # YOLOv8-Base config
+    class Yolov3BaseConfig(object):
+        def __init__(self) -> None:
+            # ---------------- Model config ----------------
+            self.width    = 1.0
+            self.depth    = 1.0
+            self.out_stride = [8, 16, 32]
+            self.max_stride = 32
+            self.num_levels = 3
+            ## Head
+            self.head_dim  = 256
+            self.anchor_size = {0: [[10, 13],   [16, 30],   [33, 23]],
+                                1: [[30, 61],   [62, 45],   [59, 119]],
+                                2: [[116, 90],  [156, 198], [373, 326]]}
+
+    cfg = Yolov3BaseConfig()
+    cfg.num_classes = 20
+    # Build a pred layer
+    pred = Yolov3DetPredLayer(cfg)
+
+    # Inference
+    cls_feats = [torch.randn(1, cfg.head_dim, 80, 80),
+                 torch.randn(1, cfg.head_dim, 40, 40),
+                 torch.randn(1, cfg.head_dim, 20, 20),]
+    reg_feats = [torch.randn(1, cfg.head_dim, 80, 80),
+                 torch.randn(1, cfg.head_dim, 40, 40),
+                 torch.randn(1, cfg.head_dim, 20, 20),]
+    t0 = time.time()
+    output = pred(cls_feats, reg_feats)
+    t1 = time.time()
+    print('Time: ', t1 - t0)
+    print('====== Pred output ======= ')
+    pred_obj = output["pred_obj"]
+    pred_cls = output["pred_cls"]
+    pred_reg = output["pred_reg"]
+    pred_box = output["pred_box"]
+    anchors  = output["anchors"]
+    
+    for level in range(cfg.num_levels):
+        print("- Level-{} : objectness       -> {}".format(level, pred_obj[level].shape))
+        print("- Level-{} : classification   -> {}".format(level, pred_cls[level].shape))
+        print("- Level-{} : delta regression -> {}".format(level, pred_reg[level].shape))
+        print("- Level-{} : bbox regression  -> {}".format(level, pred_box[level].shape))
+        print("- Level-{} : anchor boxes     -> {}".format(level, anchors[level].shape))
+
+    flops, params = profile(pred, inputs=(cls_feats, reg_feats, ), verbose=False)
+    print('==============================')
+    print('GFLOPs : {:.2f}'.format(flops / 1e9 * 2))
+    print('Params : {:.2f} M'.format(params / 1e6))