modify notes in yolov1 pred

yolo/models/yolov1/yolov1_head.py

@@ -83,11 +83,9 @@ class Yolov1DetHead(nn.Module):
 
 
 if __name__=='__main__':
-    import time
     from thop import profile
-    # Model config
     
-    # YOLOv8-Base config
+    # YOLOv1 configuration
     class Yolov1BaseConfig(object):
         def __init__(self) -> None:
             # ---------------- Model config ----------------
@@ -98,23 +96,24 @@ if __name__=='__main__':
             self.head_norm = 'BN'
             self.head_depthwise = False
             self.head_dim  = 256
-            self.num_cls_head   = 2
-            self.num_reg_head   = 2
-
+            self.num_cls_head = 2
+            self.num_reg_head = 2
     cfg = Yolov1BaseConfig()
+
     # Build a head
-    head = Yolov1DetHead(cfg, 512)
+    model = Yolov1DetHead(cfg, 512)
 
+    # Randomly generate a input data
+    x = torch.randn(2, 512, 20, 20)
 
     # Inference
-    x = torch.randn(1, 512, 20, 20)
-    t0 = time.time()
-    cls_feat, reg_feat = head(x)
-    t1 = time.time()
-    print('Time: ', t1 - t0)
-    print(cls_feat.shape, reg_feat.shape)
+    cls_feats, reg_feats = model(x)
+    print(' - the shape of input :  ', x.shape)
+    print(' - the shape of cls feats : ', cls_feats.shape)
+    print(' - the shape of reg feats : ', reg_feats.shape)
 
-    flops, params = profile(head, inputs=(x, ), verbose=False)
-    print('==============================')
-    print('GFLOPs : {:.2f}'.format(flops / 1e9 * 2))
-    print('Params : {:.2f} M'.format(params / 1e6))    
+    x = torch.randn(1, 512, 20, 20)
+    flops, params = profile(model, inputs=(x, ), verbose=False)
+    print('============== FLOPs & Params ================')
+    print(' - FLOPs  : {:.2f} G'.format(flops / 1e9 * 2))
+    print(' - Params : {:.2f} M'.format(params / 1e6))

yolo/models/yolov1/yolov1_neck.py

@@ -48,9 +48,7 @@ class SPPF(nn.Module):
 
 
 if __name__=='__main__':
-    import time
     from thop import profile
-    # Model config
     
     # YOLOv1 configuration
     class Yolov1BaseConfig(object):

yolo/models/yolov1/yolov1_pred.py

@@ -96,11 +96,10 @@ class Yolov1DetPredLayer(nn.Module):
 
 
 if __name__=='__main__':
-    import time
     from thop import profile
     # Model config
     
-    # YOLOv8-Base config
+    # YOLOv1 configuration
     class Yolov1BaseConfig(object):
         def __init__(self) -> None:
             # ---------------- Model config ----------------
@@ -108,19 +107,19 @@ if __name__=='__main__':
             self.max_stride = 32
             ## Head
             self.head_dim  = 512
-
     cfg = Yolov1BaseConfig()
     cfg.num_classes = 20
+
     # Build a pred layer
-    pred = Yolov1DetPredLayer(cfg)
+    model = Yolov1DetPredLayer(cfg)
+
+    # Randomly generate a input data
+    cls_feat = torch.randn(2, cfg.head_dim, 20, 20)
+    reg_feat = torch.randn(2, cfg.head_dim, 20, 20)
 
     # 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)
+    output = model(cls_feat, reg_feat)
+
     print('====== Pred output ======= ')
     for k in output:
         if isinstance(output[k], torch.Tensor):
@@ -128,7 +127,9 @@ if __name__=='__main__':
         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))
+    cls_feat = torch.randn(1, cfg.head_dim, 20, 20)
+    reg_feat = torch.randn(1, cfg.head_dim, 20, 20)
+    flops, params = profile(model, inputs=(cls_feat, reg_feat, ), verbose=False)
+    print('============== FLOPs & Params ================')
+    print(' - FLOPs  : {:.2f} G'.format(flops / 1e9 * 2))
+    print(' - Params : {:.2f} M'.format(params / 1e6))