Unable to Run Inference with Provided Checkpoint (imdy_pretrain.pt)

[OkayJeff5]

Hi team,

First of all, thank you for sharing this great work.

I'm currently trying to run inference using the imdy_pretrain.pt checkpoint provided in the repository. I followed the expected procedure by initializing the model from models/IDFD_mkr.py and loading the config file from config/IDFD_mkr.yml.

However, I encountered several issues during the process:

1. Missing Config Keys

The config file is missing several required keys under the transformer section, such as:

• dim
• num_head
• num_layers
• norm

I was able to resolve this by adding them manually based on how they're used in the model definition, but I am not sure if the values I used match the training configuration.

2. Shape Mismatch when Loading Weights

When calling:

model.load_state_dict(ckpt['model_state'])

I got the following error:

RuntimeError: Error(s) in loading state_dict for mkrIDFD:
    size mismatch for ID_inProj.weight: copying a param with shape torch.Size([64, 36]) from checkpoint, the shape in current model is torch.Size([64, 72]).
    size mismatch for ID_outProj.weight: copying a param with shape torch.Size([332, 64]) from checkpoint, the shape in current model is torch.Size([236, 64]).
    size mismatch for ID_outProj.bias: copying a param with shape torch.Size([332]) from checkpoint, the shape in current model is torch.Size([236]).
    size mismatch for FD_inProj.mkr_pre.weight: copying a param with shape torch.Size([64, 9]) from checkpoint, the shape in current model is torch.Size([64, 18]).

This suggests that either:

• The model structure changed since the checkpoint was created, or

• The config I used does not match the one used during training.

Could you please share the exact config file (YAML or other format) used to train the imdy_pretrain.pt checkpoint?

Additionally, a minimal working inference script or example input/output would be extremely helpful to ensure proper usage of the model and checkpoint.

Thanks again for your contribution to the community, and I appreciate your help!

[hitefork]

Hi! Sorry for our late response. We appreciate you engaging with our work.

1. Missing Config Keys

The required transformer parameters (dim, num_head, num_layers, norm) have been defined in config/IDFD_mkr.yml:

ImDy/config/IDFD_mkr.yml

Line 58 in 758820c

norm: 'layer'

ImDy/config/IDFD_mkr.yml

Line 59 in 758820c

num_layers: 3

ImDy/config/IDFD_mkr.yml

Line 60 in 758820c

dim: 64

ImDy/config/IDFD_mkr.yml

Line 61 in 758820c

num_head: 4

2.Shape Mismatch when Loading Weights

This occurs when the model configuration doesn't match the checkpoint's expected structure. If you follow our config correctly, I thing your issue can be solved.

Please let us know if you need further assistance!