Push models to hub with custom embed_layer

[Laurent2916]

Hi !

Let's say I have the following ViT model, which uses a custom embed_layer:

model = timm.create_model('vit_tiny_patch16_224', embed_layer=MyEmbedLayer)

How is it possible to push it to the hub (using timm.models.push_to_hf_hub) such that users may download it and use it without having to copy-paste MyEmbedLayer's implementation and without specifying embed_layer=MyEmbedLayer in create_model ?

For example, doing the following:

timm.models.push_to_hf_hub(
    model,
    repo_id="vit_tiny_patch16_224.embed_layer_custom_MyEmbedLayer",
    private=True,
)

downloaded_model = timm.create_model(
    'hf-hub:1aurent/vit_tiny_patch16_224.embed_layer_custom_MyEmbedLayer',
    pretrained=True
)

results in a error because the weights cannot be mapped properly:

RuntimeError: Error(s) in loading state_dict for VisionTransformer:
	Missing key(s) in state_dict: "patch_embed.proj.weight", "patch_embed.proj.bias". 
	Unexpected key(s) in state_dict: "patch_embed.proj.0.weight", "patch_embed.proj.1.weight", "patch_embed.proj.1.bias", "patch_embed.proj.1.running_mean", "patch_embed.proj.1.running_var", "patch_embed.proj.1.num_batches_tracked", "patch_embed.proj.3.weight", "patch_embed.proj.4.weight", "patch_embed.proj.4.bias", "patch_embed.proj.4.running_mean", "patch_embed.proj.4.running_var", "patch_embed.proj.4.num_batches_tracked", "patch_embed.proj.6.weight", "patch_embed.proj.6.bias". 

This can obviously can be fixed by doing:

downloaded_model = timm.create_model(
    'hf-hub:1aurent/vit_tiny_patch16_224.embed_layer_custom_MyEmbedLayer',
    pretrained=True,
    embed_layer=MyEmbedLayer
)

But this requires the users to copy-paste the MyEmbedLayer's implementation.

I know that it's possible to push model with custom_code on the hub, but this requires the transformers library.
I was wondering if this was possible in "pure timm".

Thanks !

[rwightman]

@Laurent2916 yes, I don't currently have a means of supporting custom code via the hub like transformers does. I actually havevn't looked into it super closely. If you've dug in and have any ideas / feel it's possible, I'd be open to working on a PR. It'd be useful for this case, and also for supporting 'adapters' that add non-trivial heads to the models, other custom layers that can be configured, etc.

I noticed you were uploading some neat timm (& transformers) models to the HF hub, nice!

[Laurent2916]

Hi @rwightman !

I see, well I'm not familiar enough with the code yet to suggest something that would work well. I'll try to look a bit into it though!

Thanks for noticing the couple models I've uploaded on the hub. Digging up obscure GitHub repos (which sometimes feels like archaeology) to upload model weights to the hub is actually a nice way to kill time while looking for a job 🥲.