🔧 [Add] model name into modelcfg, and move reg_max

yolo/config/config.py

@@ -24,6 +24,7 @@ class BlockConfig:
 
 @dataclass
 class ModelConfig:
+    name: Optional[str]
     anchor: AnchorConfig
     model: Dict[str, BlockConfig]
 

yolo/config/model/v9-c.yaml

@@ -1,3 +1,5 @@
+name: v9-c
+
 anchor:
   reg_max: 16
   strides: [8, 16, 32]
@@ -73,8 +75,6 @@ model:
     - MultiheadDetection:
         source: [P3, P4, P5]
         tags: Main
-        args:
-            reg_max: ${model.anchor.reg_max}
         output: True
 
   auxiliary:
@@ -129,6 +129,4 @@ model:
     - MultiheadDetection:
         source: [A3, A4, A5]
         tags: AUX
-        args:
-            reg_max: ${model.anchor.reg_max}
         output: True

yolo/config/model/v9-m.yaml

@@ -1,3 +1,5 @@
+name: v9-m
+
 anchor:
   reg_max: 16
 
@@ -72,8 +74,6 @@ model:
     - MultiheadDetection:
         source: [P3, P4, P5]
         tags: Main
-        args:
-            reg_max: ${model.anchor.reg_max}
         output: True
 
   auxiliary:
@@ -128,6 +128,4 @@ model:
     - MultiheadDetection:
         source: [A3, A4, A5]
         tags: AUX
-        args:
-            reg_max: ${model.anchor.reg_max}
         output: True

yolo/config/model/v9-s.yaml

@@ -1,3 +1,5 @@
+name: v9-s
+
 anchor:
   reg_max: 16
 
@@ -92,8 +94,6 @@ model:
     - MultiheadDetection:
         source: [P3, P4, P5]
         tags: Main
-        args:
-            reg_max: ${model.anchor.reg_max}
         output: True
 
   auxiliary:
@@ -129,6 +129,4 @@ model:
     - MultiheadDetection:
         source: [A3, A4, A5]
         tags: AUX
-        args:
-            reg_max: ${model.anchor.reg_max}
         output: True

yolo/model/yolo.py

@@ -25,8 +25,9 @@ class YOLO(nn.Module):
         self.num_classes = class_num
         self.layer_map = get_layer_map()  # Get the map Dict[str: Module]
         self.model: List[YOLOLayer] = nn.ModuleList()
-        self.build_model(model_cfg.model)
+        self.reg_max = getattr(model_cfg.anchor, "reg_max", 16)
         self.strides = getattr(model_cfg.anchor, "strides", None)
+        self.build_model(model_cfg.model)
 
     def build_model(self, model_arch: Dict[str, List[Dict[str, Dict[str, Dict]]]]):
         self.layer_index = {}
@@ -48,6 +49,7 @@ class YOLO(nn.Module):
                 if "Detection" in layer_type:
                     layer_args["in_channels"] = [output_dim[idx] for idx in source]
                     layer_args["num_classes"] = self.num_classes
+                    layer_args["reg_max"] = self.reg_max
 
                 # create layers
                 layer = self.create_layer(layer_type, source, layer_info, **layer_args)