yjh0410 9 月之前
父節點
當前提交
f82ac868a4
共有 2 個文件被更改,包括 22 次插入6 次删除
  1. 1 1
      yolo/engine.py
  2. 21 5
      yolo/models/yolov7/yolov7_backbone.py

+ 1 - 1
yolo/engine.py

@@ -163,7 +163,7 @@ class YoloTrainer(object):
         metric_logger.add_meter('gnorm', SmoothedValue(window_size=1, fmt='{value:.1f}'))
         header = 'Epoch: [{} / {}]'.format(self.epoch, self.cfg.max_epoch)
         epoch_size = len(self.train_loader)
-        print_freq = 10
+        print_freq = 100
         gnorm = 0.0
 
         # basic parameters

+ 21 - 5
yolo/models/yolov7/yolov7_backbone.py

@@ -11,7 +11,7 @@ in1k_pretrained_urls = {
     "elannet_large": "https://github.com/yjh0410/image_classification_pytorch/releases/download/weight/yolov7_elannet_large.pth",
 }
 
-# --------------------- Yolov7 backbone (CSPDarkNet-53 with SiLU) -----------------------
+# --------------------- Yolov7 backbone -----------------------
 class Yolov7Backbone(nn.Module):
     def __init__(self, use_pretrained: bool = False):
         super(Yolov7Backbone, self).__init__()
@@ -30,22 +30,38 @@ class Yolov7Backbone(nn.Module):
         ## P2/4: Stage-1
         self.layer_2 = nn.Sequential(   
             ConvModule(self.feat_dims[1], self.feat_dims[2], kernel_size=3, stride=2),             
-            ELANBlock(self.feat_dims[2], self.feat_dims[3], self.squeeze_ratios[0], self.branch_depths[0])
+            ELANBlock(in_dim = self.feat_dims[2],
+                      out_dim = self.feat_dims[3],
+                      expansion = self.squeeze_ratios[0],
+                      branch_depth = self.branch_depths[0],
+                      )
         )
         ## P3/8: Stage-2
         self.layer_3 = nn.Sequential(
             DownSample(self.feat_dims[3], self.feat_dims[3]),
-            ELANBlock(self.feat_dims[3], self.feat_dims[4], self.squeeze_ratios[1], self.branch_depths[1])
+            ELANBlock(in_dim = self.feat_dims[3],
+                      out_dim = self.feat_dims[4],
+                      expansion = self.squeeze_ratios[1],
+                      branch_depth = self.branch_depths[1],
+                      )
         )
         ## P4/16: Stage-3
         self.layer_4 = nn.Sequential(
             DownSample(self.feat_dims[4], self.feat_dims[4]),
-            ELANBlock(self.feat_dims[4], self.feat_dims[5], self.squeeze_ratios[2], self.branch_depths[2])
+            ELANBlock(in_dim = self.feat_dims[4],
+                      out_dim = self.feat_dims[5],
+                      expansion = self.squeeze_ratios[2],
+                      branch_depth = self.branch_depths[2],
+                      )
         )
         ## P5/32: Stage-4
         self.layer_5 = nn.Sequential(
             DownSample(self.feat_dims[5], self.feat_dims[5]),
-            ELANBlock(self.feat_dims[5], self.feat_dims[6], self.squeeze_ratios[3], self.branch_depths[3])
+            ELANBlock(in_dim = self.feat_dims[5],
+                      out_dim = self.feat_dims[6],
+                      expansion = self.squeeze_ratios[3],
+                      branch_depth = self.branch_depths[3],
+                      )
         )
 
         # Initialize all layers