🐛 [Fix] some bugs, fit the create_model, device

yolo/lazy.py

@@ -22,7 +22,7 @@ def main(cfg: Config):
     dataloader = create_dataloader(cfg.task.data, cfg.dataset, cfg.task.task)
     device = torch.device(cfg.device)
     if getattr(cfg.task, "fast_inference", False):
-        model = FastModelLoader(cfg).load_model()
+        model = FastModelLoader(cfg, device).load_model()
         device = torch.device(cfg.device)
     else:
         model = create_model(cfg.model, class_num=cfg.class_num, weight_path=cfg.weight, device=device)

yolo/model/yolo.py

@@ -43,7 +43,7 @@ class YOLO(nn.Module):
                 source = self.get_source_idx(layer_info.get("source", -1), layer_idx)
 
                 # Find in channels
-                if any(module in layer_type for module in ["Conv", "ELAN", "ADown", "CBLinear"]):
+                if any(module in layer_type for module in ["Conv", "ELAN", "ADown", "AConv", "CBLinear"]):
                     layer_args["in_channels"] = output_dim[source]
                 if "Detection" in layer_type:
                     layer_args["in_channels"] = [output_dim[idx] for idx in source]
@@ -81,7 +81,7 @@ class YOLO(nn.Module):
         return output
 
     def get_out_channels(self, layer_type: str, layer_args: dict, output_dim: list, source: Union[int, list]):
-        if any(module in layer_type for module in ["Conv", "ELAN", "ADown"]):
+        if any(module in layer_type for module in ["Conv", "ELAN", "ADown", "AConv"]):
             return layer_args["out_channels"]
         if layer_type == "CBFuse":
             return output_dim[source[-1]]
@@ -134,7 +134,7 @@ def create_model(model_cfg: ModelConfig, weight_path: Optional[str], device: dev
             logger.info(f"🌐 Weight {weight_path} not found, try downloading")
             prepare_weight(weight_path=weight_path)
         if os.path.exists(weight_path):
-            model.model.load_state_dict(torch.load(weight_path, map_location=device))
+            model.model.load_state_dict(torch.load(weight_path, map_location=device), strict=False)
             logger.info("✅ Success load model weight")
 
     log_model_structure(model.model)

yolo/tools/solver.py

@@ -143,8 +143,7 @@ class ModelTester:
                         break
                     if not self.save_predict:
                         continue
-
-                if self.save_predict == False:
+                if self.save_predict != False:
                     save_image_path = os.path.join(self.save_path, f"frame{idx:03d}.png")
                     img.save(save_image_path)
                     logger.info(f"💾 Saved visualize image at {save_image_path}")