Keys not found during inference

[fronos]

Describe the bug

I'm trained custom model Yolo v9-m from scratch. During inference I noticed that logger throwing messages:

WARNING  ⚠ Weight Not Found for key: 19.conv                                                                                                                                      yolo.py:152
WARNING  ⚠ Weight Not Found for key: 21.conv3.0.conv2                                                                                                                             yolo.py:152
WARNING  ⚠ Weight Not Found for key: 34.conv3.0.conv2                                                                                                                             yolo.py:152
WARNING  ⚠ Weight Not Found for key: 38.heads.1.anchor_conv.0                                                                                                                     yolo.py:152
WARNING  ⚠ Weight Not Found for key: 22.heads.2.anc2vec  

The main reasons of this warnings:

1. Lightning save weights as a key of file with extension .ckpt.
2. Keys of initialized model and trained model mismatched. Need add model.model into name of key

System Info (please complete the following ## information):

• OS: [e.g. Ubuntu 20.04]
• Python Version: 3.10.16
• PyTorch Version: 2.5.1
• Lightning: 2.4.0
• CUDA/cuDNN/MPS Version: 12.4
• YOLO Model Version: YOLOv9-m

[cansik]

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")