Problem with trex_inference.py 'size mismatch for encoder.sentence_encoder.byte_combine.projection.weight'

[caroline-mrs]

I'm unable to execute the trex_inference.py file as it's returning this error:

[D] Loading function traces
[D] Loading Trex model
Traceback (most recent call last):
  File "trex_inference.py", line 102, in <module>
    main(input_pairs, input_traces, model_checkpoint_dir,data_bin_dir, output_dir)
  File "trex_inference.py", line 60, in main
    trex = TrexModel.from_pretrained(f'checkpoints/similarity',
  File "/fairseq/models/trex/model.py", line 308, in from_pretrained
    x = hub_utils.from_pretrained(
  File "/fairseq/hub_utils.py", line 73, in from_pretrained
    models, args, task = checkpoint_utils.load_model_ensemble_and_task(
  File "/fairseq/checkpoint_utils.py", line 389, in load_model_ensemble_and_task
    model.load_state_dict(state["model"], strict=strict, model_cfg=cfg.model)
  File "/fairseq/models/fairseq_model.py", line 125, in load_state_dict
    return super().load_state_dict(new_state_dict, strict)
  File "/opt/conda/envs/trex/lib/python3.8/site-packages/torch/nn/modules/module.py", line 2189, in load_state_dict
    raise RuntimeError('Error(s) in loading state_dict for {}:\n\t{}'.format(
RuntimeError: Error(s) in loading state_dict for TrexModel:
        Unexpected key(s) in state_dict: "encoder.sentence_encoder.embed_bytes.weight". 
        size mismatch for encoder.sentence_encoder.byte_combine.convolutions.0.weight: copying a param with shape torch.Size([64, 768, 1]) from checkpoint, the shape in current model is torch.Size([4, 1, 1]).
        size mismatch for encoder.sentence_encoder.byte_combine.convolutions.0.bias: copying a param with shape torch.Size([64]) from checkpoint, the shape in current model is torch.Size([4]).
        size mismatch for encoder.sentence_encoder.byte_combine.convolutions.1.weight: copying a param with shape torch.Size([128, 768, 2]) from checkpoint, the shape in current model is torch.Size([8, 1, 2]).
        size mismatch for encoder.sentence_encoder.byte_combine.convolutions.1.bias: copying a param with shape torch.Size([128]) from checkpoint, the shape in current model is torch.Size([8]).
        size mismatch for encoder.sentence_encoder.byte_combine.convolutions.2.weight: copying a param with shape torch.Size([192, 768, 3]) from checkpoint, the shape in current model is torch.Size([12, 1, 3]).
        size mismatch for encoder.sentence_encoder.byte_combine.convolutions.2.bias: copying a param with shape torch.Size([192]) from checkpoint, the shape in current model is torch.Size([12]).
        size mismatch for encoder.sentence_encoder.byte_combine.highway.layers.0.weight: copying a param with shape torch.Size([768, 384]) from checkpoint, the shape in current model is torch.Size([48, 24]).
        size mismatch for encoder.sentence_encoder.byte_combine.highway.layers.0.bias: copying a param with shape torch.Size([768]) from checkpoint, the shape in current model is torch.Size([48]).
        size mismatch for encoder.sentence_encoder.byte_combine.highway.layers.1.weight: copying a param with shape torch.Size([768, 384]) from checkpoint, the shape in current model is torch.Size([48, 24]).
        size mismatch for encoder.sentence_encoder.byte_combine.highway.layers.1.bias: copying a param with shape torch.Size([768]) from checkpoint, the shape in current model is torch.Size([48]).
        size mismatch for encoder.sentence_encoder.byte_combine.projection.weight: copying a param with shape torch.Size([768, 384]) from checkpoint, the shape in current model is torch.Size([768, 24]).

Could someone please assist me in resolving this issue? I would be very appreciate.

[jimmy-sonny]

Hi @caroline-mrs , thank you for opening this issue.

I cannot reproduce the error that you are facing, can you please try again with the updated Trex Dockerfile - related to Issue #56?

Are you using the docker container as described in the Trex README?

Can you please provide the detailed steps and instructions that you are following?