🐛 [Fix] a bug when loading the model

yolo/config/general.yaml

@@ -12,4 +12,4 @@ lucky_number: 10
 use_wandb: False
 use_TensorBoard: False
 
-weight: weights/v9-c.pt
+weight: True # Path to weight or True for auto, False for no pretrained weight

yolo/lazy.py

@@ -22,10 +22,10 @@ def main(cfg: Config):
     device, use_ddp = get_device(cfg.device)
     dataloader = create_dataloader(cfg.task.data, cfg.dataset, cfg.task.task, use_ddp)
     if getattr(cfg.task, "fast_inference", False):
-        model = FastModelLoader(cfg).load_model()
+        model = FastModelLoader(cfg).load_model(device)
     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/yolo.py

@@ -119,7 +119,7 @@ class YOLO(nn.Module):
             raise ValueError(f"Unsupported layer type: {layer_type}")
 
 
-def create_model(model_cfg: ModelConfig, weight_path: Optional[str], class_num: int = 80) -> YOLO:
+def create_model(model_cfg: ModelConfig, weight_path: Union[bool, str], class_num: int = 80) -> YOLO:
     """Constructs and returns a model from a Dictionary configuration file.
 
     Args:
@@ -132,12 +132,14 @@ def create_model(model_cfg: ModelConfig, weight_path: Optional[str], class_num:
     OmegaConf.set_struct(model_cfg, False)
     model = YOLO(model_cfg, class_num)
     if weight_path:
+        if weight_path == True:
+            weight_path = os.path.join("weights", f"{model_cfg.name}.pt")
         if not os.path.exists(weight_path):
             logger.info(f"🌐 Weight {weight_path} not found, try downloading")
             prepare_weight(weight_path=weight_path)
         if os.path.exists(weight_path):
             # TODO: fix map_location
-            model.model.load_state_dict(torch.load(weight_path), strict=False)
+            model.model.load_state_dict(torch.load(weight_path, map_location=torch.device("cpu")), strict=False)
             logger.info("✅ Success load model & weight")
     else:
         logger.info("✅ Success load model")

yolo/utils/deploy_utils.py

@@ -13,6 +13,8 @@ class FastModelLoader:
         self.cfg = cfg
         self.compiler = cfg.task.fast_inference
         self._validate_compiler()
+        if cfg.weight == True:
+            cfg.weight = os.path.join("weights", f"{cfg.model.name}.pt")
         self.model_path = f"{os.path.splitext(cfg.weight)[0]}.{self.compiler}"
 
     def _validate_compiler(self):
@@ -23,14 +25,14 @@ class FastModelLoader:
             logger.warning("🍎 TensorRT does not support MPS devices. Using original model.")
             self.compiler = None
 
-    def load_model(self):
+    def load_model(self, device):
         if self.compiler == "onnx":
             return self._load_onnx_model()
         elif self.compiler == "trt":
-            return self._load_trt_model()
+            return self._load_trt_model().to(device)
         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)
+        return create_model(self.cfg.model, class_num=self.cfg.class_num, weight_path=self.cfg.weight).to(device)
 
     def _load_onnx_model(self):
         from onnxruntime import InferenceSession