🐛 [Update] some bug or vaiable name in Vec2Box

yolo/__init__.py

@@ -3,7 +3,7 @@ from yolo.model.yolo import create_model
 from yolo.tools.data_loader import AugmentationComposer, create_dataloader
 from yolo.tools.drawer import draw_bboxes
 from yolo.tools.solver import ModelTester, ModelTrainer, ModelValidator
-from yolo.utils.bounding_box_utils import bbox_nms
+from yolo.utils.bounding_box_utils import Vec2Box, bbox_nms
 from yolo.utils.deploy_utils import FastModelLoader
 from yolo.utils.logging_utils import custom_logger
 
@@ -13,6 +13,7 @@ all = [
     "custom_logger",
     "validate_log_directory",
     "draw_bboxes",
+    "Vec2Box",
     "bbox_nms",
     "AugmentationComposer",
     "create_dataloader",

yolo/config/model/v9-c.yaml

@@ -68,6 +68,14 @@ model:
         args: {out_channels: 512, part_channels: 512}
         tags: P5
 
+  detection:
+    - MultiheadDetection:
+        source: [P3, P4, P5]
+        tags: Main
+        args:
+            reg_max: ${model.anchor.reg_max}
+        output: True
+
   auxiliary:
     - CBLinear:
         source: B3
@@ -123,11 +131,3 @@ model:
         args:
             reg_max: ${model.anchor.reg_max}
         output: True
-
-  detection:
-    - MultiheadDetection:
-        source: [P3, P4, P5]
-        tags: Main
-        args:
-            reg_max: ${model.anchor.reg_max}
-        output: True

yolo/tools/format_converters.py

@@ -1,12 +1,13 @@
 def convert_weight(old_state_dict, new_state_dict, model_size: int = 38):
     # TODO: need to refactor
+    shift = 1
     for idx in range(model_size):
         new_list, old_list = [], []
         for weight_name, weight_value in new_state_dict.items():
             if weight_name.split(".")[0] == str(idx):
                 new_list.append((weight_name, None))
         for weight_name, weight_value in old_state_dict.items():
-            if f"model.{idx+1}." in weight_name:
+            if f"model.{idx+shift}." in weight_name:
                 old_list.append((weight_name, weight_value))
         if len(new_list) == len(old_list):
             for (weight_name, _), (_, weight_value) in zip(new_list, old_list):
@@ -17,7 +18,8 @@ def convert_weight(old_state_dict, new_state_dict, model_size: int = 38):
                     continue
                 _, _, conv_name, conv_idx, *details = weight_name.split(".")
                 if conv_name == "cv4" or conv_name == "cv5":
-                    layer_idx = 38
+                    layer_idx = 22
+                    shift = 2
                 else:
                     layer_idx = 37
 

yolo/tools/solver.py

@@ -32,7 +32,7 @@ class ModelTrainer:
         self.num_epochs = cfg.task.epoch
 
         self.validation_dataloader = create_dataloader(cfg.task.validation.data, cfg.dataset, cfg.task.validation.task)
-        self.validator = ModelValidator(cfg.task.validation, model, save_path, device, self.progress)
+        self.validator = ModelValidator(cfg.task.validation, model, vec2box, save_path, device, self.progress)
 
         if getattr(train_cfg.ema, "enabled", False):
             self.ema = ExponentialMovingAverage(model, decay=train_cfg.ema.decay)
@@ -40,14 +40,14 @@ class ModelTrainer:
             self.ema = None
         self.scaler = GradScaler()
 
-    def train_one_batch(self, data: Tensor, targets: Tensor):
-        data, targets = data.to(self.device), targets.to(self.device)
+    def train_one_batch(self, images: Tensor, targets: Tensor):
+        images, targets = images.to(self.device), targets.to(self.device)
         self.optimizer.zero_grad()
 
         with autocast():
-            outputs = self.model(data)
-            aux_predicts = self.vec2box(outputs["AUX"])
-            main_predicts = self.vec2box(outputs["Main"])
+            predicts = self.model(images)
+            aux_predicts = self.vec2box(predicts["AUX"])
+            main_predicts = self.vec2box(predicts["Main"])
             loss, loss_item = self.loss_fn(aux_predicts, main_predicts, targets)
 
         self.scaler.scale(loss).backward()
@@ -60,8 +60,8 @@ class ModelTrainer:
         self.model.train()
         total_loss = 0
 
-        for data, targets in dataloader:
-            loss, loss_each = self.train_one_batch(data, targets)
+        for images, targets in dataloader:
+            loss, loss_each = self.train_one_batch(images, targets)
 
             total_loss += loss
             self.progress.one_batch(loss_each)
@@ -111,14 +111,15 @@ class ModelTester:
 
     def solve(self, dataloader: StreamDataLoader):
         logger.info("👀 Start Inference!")
-
+        if isinstance(self.model, torch.nn.Module):
+            self.model.eval()
         try:
             for idx, images in enumerate(dataloader):
                 images = images.to(self.device)
                 with torch.no_grad():
-                    outputs = self.model(images)
-                outputs = self.vec2box(outputs["Main"])
-                nms_out = bbox_nms(outputs[0], outputs[2], self.nms)
+                    predicts = self.model(images)
+                predicts = self.vec2box(predicts["Main"])
+                nms_out = bbox_nms(predicts[0], predicts[2], self.nms)
                 draw_bboxes(
                     images[0],
                     nms_out[0],
@@ -141,15 +142,18 @@ class ModelValidator:
         self,
         validation_cfg: ValidationConfig,
         model: YOLO,
+        vec2box: Vec2Box,
         save_path: str,
         device,
         # TODO: think Progress?
         progress: ProgressTracker,
     ):
         self.model = model
+        self.vec2box = vec2box
         self.device = device
         self.progress = progress
         self.save_path = save_path
+
         self.nms = validation_cfg.nms
 
     def solve(self, dataloader):
@@ -159,11 +163,12 @@ class ModelValidator:
         iou_thresholds = torch.arange(0.5, 1.0, 0.05)
         map_all = []
         self.progress.start_one_epoch(len(dataloader))
-        for data, targets in dataloader:
-            data, targets = data.to(self.device), targets.to(self.device)
+        for images, targets in dataloader:
+            images, targets = images.to(self.device), targets.to(self.device)
             with torch.no_grad():
-                raw_output = self.model(data)
-            nms_out = bbox_nms(raw_output[-1][0], self.nms)
+                predicts = self.model(images)
+            predicts = self.vec2box(predicts["Main"])
+            nms_out = bbox_nms(predicts[0], predicts[2], self.nms)
             for idx, predict in enumerate(nms_out):
                 map_value = calculate_map(predict, targets[idx], iou_thresholds)
                 map_all.append(map_value[0])

yolo/utils/bounding_box_utils.py

@@ -4,6 +4,7 @@ from typing import List, Tuple
 import torch
 import torch.nn.functional as F
 from einops import rearrange
+from loguru import logger
 from torch import Tensor
 from torchvision.ops import batched_nms
 
@@ -264,6 +265,7 @@ class BoxMatcher:
 
 class Vec2Box:
     def __init__(self, model, image_size, device):
+        logger.info("🧸 Make a dummy test for auto-anchor size")
         dummy_input = torch.zeros(1, 3, *image_size).to(device)
         dummy_output = model(dummy_input)
         anchors_num = []