🐛 [Fix] a bug in deploying TensorRT model

yolo/config/model/v9-c.yaml

@@ -1,6 +1,6 @@
 anchor:
   reg_max: 16
-  anchors: [8, 16, 32]
+  strides: [8, 16, 32]
 
 model:
   backbone:

yolo/lazy.py

@@ -25,7 +25,7 @@ def main(cfg: Config):
         model = FastModelLoader(cfg).load_model()
     else:
         model = create_model(cfg.model, class_num=cfg.class_num, weight_path=cfg.weight)
-        model = model.to(device)
+    model = model.to(device)
 
     vec2box = Vec2Box(model, cfg.image_size, device)
 

yolo/model/module.py

@@ -105,7 +105,7 @@ class Anchor2Vec(nn.Module):
     def forward(self, anchor_x: Tensor) -> Tensor:
         anchor_x = rearrange(anchor_x, "B (P R) h w -> B R P h w", P=4)
         vector_x = anchor_x.softmax(dim=1)
-        vector_x = self.anc2vec(vector_x).squeeze(1)
+        vector_x = self.anc2vec(vector_x)[:, 0]
         return anchor_x, vector_x
 
 

yolo/model/yolo.py

@@ -26,13 +26,15 @@ class YOLO(nn.Module):
         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.strides = getattr(model_cfg.anchor, "strides", None)
 
     def build_model(self, model_arch: Dict[str, List[Dict[str, Dict[str, Dict]]]]):
         self.layer_index = {}
         output_dim, layer_idx = [3], 1
         logger.info(f"🚜 Building YOLO")
         for arch_name in model_arch:
-            logger.info(f"  🏗️  Building {arch_name}")
+            if model_arch[arch_name]:
+                logger.info(f"  🏗️  Building {arch_name}")
             for layer_idx, layer_spec in enumerate(model_arch[arch_name], start=layer_idx):
                 layer_type, layer_info = next(iter(layer_spec.items()))
                 layer_args = layer_info.get("args", {})
@@ -45,7 +47,6 @@ 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
@@ -134,6 +135,7 @@ def create_model(model_cfg: ModelConfig, weight_path: Optional[str], class_num:
         if os.path.exists(weight_path):
             # TODO: fix map_location
             model.model.load_state_dict(torch.load(weight_path), strict=False)
-            logger.info("✅ Success load model weight")
-
+            logger.info("✅ Success load model & weight")
+    else:
+        logger.info("✅ Success load model")
     return model

yolo/tools/drawer.py

@@ -13,7 +13,7 @@ def draw_bboxes(
     img: Union[Image.Image, torch.Tensor],
     bboxes: List[List[Union[int, float]]],
     *,
-    idx2label: Optional[list],
+    idx2label: Optional[list] = None,
 ):
     """
     Draw bounding boxes on an image.
@@ -47,7 +47,7 @@ def draw_bboxes(
         draw.rounded_rectangle(bbox, outline=(*color_map, 200), radius=5, width=2)
         draw.rounded_rectangle(bbox, fill=(*color_map, 100), radius=5)
 
-        class_text = str(idx2label[int(class_id)] if idx2label else class_id)
+        class_text = str(idx2label[int(class_id)] if idx2label else int(class_id))
         label_text = f"{class_text}" + (f" {conf[0]: .0%}" if conf else "")
 
         text_bbox = font.getbbox(label_text)

yolo/utils/bounding_box_utils.py

@@ -9,6 +9,7 @@ from torch import Tensor
 from torchvision.ops import batched_nms
 
 from yolo.config.config import MatcherConfig, ModelConfig, NMSConfig
+from yolo.model.yolo import YOLO
 
 
 def calculate_iou(bbox1, bbox2, metrics="iou") -> Tensor:
@@ -264,8 +265,8 @@ class BoxMatcher:
 
 
 class Vec2Box:
-    def __init__(self, model, image_size, device, anchors: list = None):
-        if anchors is None:
+    def __init__(self, model: YOLO, image_size, device):
+        if model.strides is None:
             logger.info("🧸 Found no anchor, Make a dummy test for auto-anchor size")
             dummy_input = torch.zeros(1, 3, *image_size).to(device)
             dummy_output = model(dummy_input)
@@ -274,7 +275,8 @@ class Vec2Box:
                 _, _, *anchor_num = predict_head[2].shape
                 anchors_num.append(anchor_num)
         else:
-            anchors_num = [[image_size[0] / anchor, image_size[0] / anchor] for anchor in anchors]
+            logger.info(f"🈶 Found anchor {model.strides}")
+            anchors_num = [[image_size[0] // stride, image_size[0] // stride] for stride in model.strides]
         anchor_grid, scaler = generate_anchors(image_size, anchors_num)
         self.anchor_grid, self.scaler = anchor_grid.to(device), scaler.to(device)
         self.anchor_norm = (anchor_grid / scaler[:, None])[None].to(device)

yolo/utils/deploy_utils.py

@@ -30,9 +30,7 @@ class FastModelLoader:
             return self._load_trt_model()
         elif self.compiler == "deploy":
             self.cfg.model.model.auxiliary = {}
-        return create_model(
-            self.cfg.model, class_num=self.cfg.class_num, weight_path=self.cfg.weight, device=self.device
-        )
+        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
@@ -91,9 +89,9 @@ class FastModelLoader:
         from torch2trt import torch2trt
 
         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))
+        dummy_input = torch.ones((1, 3, *self.cfg.image_size)).cuda()
         logger.info(f"♻️ Creating TensorRT model")
-        model_trt = torch2trt(model, [dummy_input])
+        model_trt = torch2trt(model.cuda(), [dummy_input])
         torch.save(model_trt.state_dict(), self.model_path)
         logger.info(f"📥 TensorRT model saved to {self.model_path}")
         return model_trt