Loading pretrained dally, mismatch between checkpoint and initialized model

Hello!

So I have a partly trained dalle model that I wanted to test but I can’t seem to load it.

My vae params were:

{'image_size': 256,
 'num_layers': 3,
 'num_tokens': 2048,
 'codebook_dim': 512,
 'hidden_dim': 256,
 'num_resnet_blocks': 2}

My dalle params were:

'dim': 512,
'vae': vae,
'num_text_tokens': 49408,
'text_seq_len': 256,
'depth': 16,
'heads': 16,
'dim_head': 64,
'attn_dropout': 0.1,
'ff_dropout': 0.1,
'reversible': True,
'attn_types': ('full', 'axial_row', 'axial_col', 'conv_like')}

When I initialize both models and try to load my state_dict I get an error in the (I believe axial attention):

RuntimeError: Error(s) in loading state_dict for DALLE:
	size mismatch for image_pos_emb.weights_0: copying a param with shape torch.Size([1, 32, 1, 512]) from checkpoint, the shape in current model is torch.Size([1, 256, 1, 512]).
	size mismatch for image_pos_emb.weights_1: copying a param with shape torch.Size([1, 1, 32, 512]) from checkpoint, the shape in current model is torch.Size([1, 1, 256, 512]).

checking manually confirms that the initialized dalle model has a mismatch with the pretrained one.

image_pos_emb.weights_0 has a dimension of torch.Size([1, 1, 256, 512]) for the initialized model image_pos_emb.weights_0 has a dimension of torch.Size([1, 1, 32, 512]) for the pretrained model

Wonder if I’m doing something wrong here? It’s quite possible I’ve totally missed something!

[lucidrains]

@TheodoreGalanos Oh no, but that model you trained was prior to a bug fix to the positional encoding 🙏 You may have to just retrain it 😦

[TheodoreGalanos]

Nevermind, turns out the mismatch was in my local computer. All fixed after updated the dalle-pytorch package!