🔧 [Update] num_classes position, move to hyper

examples/example_train.py

@@ -23,7 +23,7 @@ def main(cfg: Config):
         prepare_dataset(cfg.download)
 
     dataloader = get_dataloader(cfg)
-    model = get_model(cfg.model)
+    model = get_model(cfg)
     draw_model(model=model)
     # TODO: get_device or rank, for DDP mode
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

yolo/config/model/v7-base.yaml

@@ -1,5 +1,3 @@
-num_classes: 80
-
 anchor:
   reg_max: 16
   strides: [8, 16, 32]

yolo/config/model/v9-c.yaml

@@ -1,4 +1,6 @@
-num_classes: 80
+anchor:
+  reg_max: 16
+  strides: [8, 16, 32]
 
 model:
   backbone:
@@ -117,3 +119,4 @@ model:
 
     - MultiheadDetection:
         source: [A3, A4, A5, P3, P4, P5]
+        output: True

yolo/model/yolo.py

@@ -4,7 +4,7 @@ import torch.nn as nn
 from loguru import logger
 from omegaconf import ListConfig, OmegaConf
 
-from yolo.config.config import Model
+from yolo.config.config import Config, Model
 from yolo.tools.layer_helper import get_layer_map
 
 
@@ -17,9 +17,9 @@ class YOLO(nn.Module):
                    parameters, and any other relevant configuration details.
     """
 
-    def __init__(self, model_cfg: Model):
+    def __init__(self, model_cfg: Model, num_classes: int):
         super(YOLO, self).__init__()
-        self.num_classes = model_cfg["num_classes"]
+        self.num_classes = num_classes
         self.layer_map = get_layer_map()  # Get the map Dict[str: Module]
         self.build_model(model_cfg.model)
 
@@ -101,7 +101,7 @@ class YOLO(nn.Module):
             raise ValueError(f"Unsupported layer type: {layer_type}")
 
 
-def get_model(model_cfg: dict) -> YOLO:
+def get_model(cfg: Config) -> YOLO:
     """Constructs and returns a model from a Dictionary configuration file.
 
     Args:
@@ -110,7 +110,7 @@ def get_model(model_cfg: dict) -> YOLO:
     Returns:
         YOLO: An instance of the model defined by the given configuration.
     """
-    OmegaConf.set_struct(model_cfg, False)
-    model = YOLO(model_cfg)
+    OmegaConf.set_struct(cfg.model, False)
+    model = YOLO(cfg.model, cfg.hyper.data.class_num)
     logger.info("✅ Success load model")
     return model