🔧 [Move] model class num config out of modelyaml

yolo/config/config.py

@@ -25,7 +25,6 @@ class BlockConfig:
 @dataclass
 class ModelConfig:
     anchor: AnchorConfig
-    class_num: int
     model: Dict[str, BlockConfig]
 
 

yolo/config/model/v9-c.yaml

@@ -2,8 +2,6 @@ anchor:
   reg_max: 16
   strides: [8, 16, 32]
 
-class_num: ${class_num}
-
 model:
   backbone:
     - Conv:

yolo/lazy.py

@@ -25,7 +25,7 @@ def main(cfg: Config):
         model = FastModelLoader(cfg).load_model()
         device = torch.device(cfg.device)
     else:
-        model = create_model(cfg.model, cfg.weight).to(device)
+        model = create_model(cfg.model, class_num=cfg.class_num, weight_path=cfg.weight).to(device)
 
     if cfg.task.task == "train":
         trainer = ModelTrainer(cfg, model, save_path, device)

yolo/model/module.py

@@ -93,13 +93,13 @@ class MultiheadDetection(nn.Module):
 
 
 class Anchor2Box(nn.Module):
-    def __init__(self, reg_max, strides) -> None:
+    def __init__(self, reg_max, strides, num_classes: int) -> None:
         super().__init__()
         self.reg_max = reg_max
         self.strides = strides
         # TODO: read by cfg!
         image_size = [640, 640]
-        self.class_num = 80
+        self.num_classes = num_classes
         self.anchors, self.scaler = generate_anchors(image_size, self.strides)
         reverse_reg = torch.arange(self.reg_max, dtype=torch.float32)
         self.reverse_reg = nn.Parameter(reverse_reg, requires_grad=False)
@@ -117,7 +117,7 @@ class Anchor2Box(nn.Module):
         for pred in predicts:
             preds.append(rearrange(pred, "B AC h w -> B (h w) AC"))  # B x AC x h x w-> B x hw x AC
         preds = torch.concat(preds, dim=1)  # -> B x (H W) x AC
-        preds_anc, preds_cls = torch.split(preds, (self.reg_max * 4, self.class_num), dim=-1)
+        preds_anc, preds_cls = torch.split(preds, (self.reg_max * 4, self.num_classes), dim=-1)
         preds_anc = rearrange(preds_anc, "B  hw (P R)-> B hw P R", P=4)
 
         pred_LTRB = preds_anc.softmax(dim=-1) @ self.reverse_reg * self.scaler.view(1, -1, 1)

yolo/model/yolo.py

@@ -22,9 +22,9 @@ class YOLO(nn.Module):
                    parameters, and any other relevant configuration details.
     """
 
-    def __init__(self, model_cfg: ModelConfig):
+    def __init__(self, model_cfg: ModelConfig, class_num: int = 80):
         super(YOLO, self).__init__()
-        self.num_classes = model_cfg.class_num
+        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)
@@ -47,6 +47,7 @@ class YOLO(nn.Module):
                     layer_args["in_channels"] = output_dim[source]
                 if "Detection" in layer_type:
                     layer_args["in_channels"] = [output_dim[idx] for idx in source]
+                if "Detection" in layer_type or "Anchor2Box" in layer_type:
                     layer_args["num_classes"] = self.num_classes
 
                 # create layers
@@ -116,7 +117,7 @@ class YOLO(nn.Module):
             raise ValueError(f"Unsupported layer type: {layer_type}")
 
 
-def create_model(model_cfg: ModelConfig, weight_path: str) -> YOLO:
+def create_model(model_cfg: ModelConfig, class_num: int = 80, weight_path: str = "weights/v9-c.pt") -> YOLO:
     """Constructs and returns a model from a Dictionary configuration file.
 
     Args:
@@ -126,7 +127,7 @@ def create_model(model_cfg: ModelConfig, weight_path: str) -> YOLO:
         YOLO: An instance of the model defined by the given configuration.
     """
     OmegaConf.set_struct(model_cfg, False)
-    model = YOLO(model_cfg)
+    model = YOLO(model_cfg, class_num)
     logger.info("✅ Success load model")
     if weight_path:
         if os.path.exists(weight_path):

yolo/tools/loss_functions.py

@@ -70,7 +70,7 @@ class DFLoss(nn.Module):
 class YOLOLoss:
     def __init__(self, cfg: Config) -> None:
         self.reg_max = cfg.model.anchor.reg_max
-        self.class_num = cfg.model.class_num
+        self.class_num = cfg.class_num
         self.image_size = list(cfg.image_size)
         self.strides = cfg.model.anchor.strides
         device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

yolo/utils/deploy_utils.py

@@ -28,7 +28,7 @@ class FastModelLoader:
             return self._load_onnx_model()
         elif self.compiler == "trt":
             return self._load_trt_model()
-        return create_model(self.cfg)
+        return create_model(self.cfg.model, class_num=self.cfg.class_num, weight_path=self.cfg.weight)
 
     def _load_onnx_model(self):
         from onnxruntime import InferenceSession
@@ -53,7 +53,7 @@ class FastModelLoader:
         from onnxruntime import InferenceSession
         from torch.onnx import export
 
-        model = create_model(self.cfg).eval()
+        model = create_model(self.cfg.model, class_num=self.cfg.class_num, weight_path=self.cfg.weight).eval()
         dummy_input = torch.ones((1, 3, *self.cfg.image_size))
         export(
             model,
@@ -81,7 +81,7 @@ class FastModelLoader:
     def _create_trt_model(self):
         from torch2trt import torch2trt
 
-        model = create_model(self.cfg).eval()
+        model = create_model(self.cfg.model, class_num=self.cfg.class_num, weight_path=self.cfg.weight).eval()
         dummy_input = torch.ones((1, 3, *self.cfg.image_size))
         logger.info(f"♻️ Creating TensorRT model")
         model_trt = torch2trt(model, [dummy_input])