Keys not found during inference #136

Activity

1. fronos

   added

   bugSomething isn't workingSomething isn't working

   on Dec 10, 2024
2. fronos

   changed the title [-]Missed keys during inference[/-] [+]Keys not found during inference[/+] on Dec 10, 2024
3. cansik commented on Dec 11, 2024

   

   cansik

   on Dec 11, 2024

   Contributor

   More actions

   I have the same issue here for a custom YOLOv9-t model. Basically, the custom trained checkpoint won't be able to load for inference. What I've done is adjusting the create_model method in yolo.py to allow to read ckpt files. This seems to work!

   def create_model(model_cfg: ModelConfig, weight_path: Union[bool, Path] = True, class_num: int = 80) -> YOLO:
       """Constructs and returns a model from a Dictionary configuration file.
   
       Args:
           config_file (dict): The configuration file of the model.
   
       Returns:
           YOLO: An instance of the model defined by the given configuration.
       """
       OmegaConf.set_struct(model_cfg, False)
       model = YOLO(model_cfg, class_num)
       if weight_path:
           if weight_path == True:
               weight_path = Path("weights") / f"{model_cfg.name}.pt"
           elif isinstance(weight_path, str):
               weight_path = Path(weight_path)
   
           if not weight_path.exists():
               logger.info(f"🌐 Weight {weight_path} not found, try downloading")
               prepare_weight(weight_path=weight_path)
           if weight_path.exists():
               if weight_path.suffix == ".ckpt":
                   checkpoint = torch.load(weight_path)
                   state_dict = checkpoint['state_dict']
   
                   # Fix the keys in the checkpoint
                   new_state_dict = {}
   
                   for key, value in state_dict.items():
                       if key.startswith("model.model."):
                           new_key = key.replace("model.model.", "model.", 1)
                           new_state_dict[new_key] = value
                       else:
                           new_state_dict[key] = value
   
                   # Load new state dict
                   model.load_state_dict(new_state_dict, strict=False)
   
                   for name, param in model.named_parameters():
                       if name not in new_state_dict:
                           print(f"Missing weight for layer: {name}")
   
                   logger.info(":white_check_mark: Success load model & checkpoint")
               else:
                   model.save_load_weights(weight_path)
                   logger.info(":white_check_mark: Success load model & weight")

   👍React with 👍3Reacted by fangliang425, NPetsky and GilSeamas
4. divyanidhi commented on Mar 28, 2025

   

   divyanidhi

   on Mar 28, 2025

   More actions

   @cansik Can you please share your run command for inferencing?

   👍React with 👍1Reacted by Sam Berkson