How to write config file to build custom ViT via HuggingFace Hub

[gau-nernst]

I'm re-hosting AudioMAE weights on HF hub so that I can load AudioMAE model easily from timm: https://huggingface.co/gaunernst/vit_base_patch16_1024_128.audiomae_as2m_ft_as20k

The ideal outcome is to be able to load the model like any other timm models:

model = timm.create_model("hf_hub:gaunernst/vit_base_patch16_1024_128.audiomae_as2m_ft_as20k", pretrained=True)

However, I'm facing some problems writing the config.json file so that the model configuration is correct. In essence, AudioMAE ViT-B has the following differences from the usual ViT:

• img_size=(1024,128)
• in_chans=1
• global_pool='avg'

After mimicking your other ViT config files and trying to change the config, I figured that setting input_size=[1, 1024, 128] would fix img_size and in_chans. However, global_pool='avg' in config.json doesn't seem to take effect, and I have to override that option in timm.create_model().

timm.create_model() code is pretty complicated, and I hope asking here should give me a faster answer. The question is, how would I specify global_pool='avg' in a HF-hosted model? And more generally, is there a mechanism to configure all model hyper-parameters for a HF-hosted model i.e. call VisionTransformer(img_size=(1024, 128), in_chans=1, global_pool='avg') directly instead of using a pre-defined model function vit_base_patch16_224() and override some options?

Thank you!

[rwightman]

@gau-nernst correct re the input_size, that is adjustable in that fashion, the global pool and other args passed through the model are not right now....

BUT, it's a completely valid use case and I would like to support this. I think I could make a model_kwargs key at the root level and pass those through on model construction. Have to tinker with that, but might be a quick add to support this use case and also allow other changes (based on model), like changing act/norm layers, changing model layer sizes, etc.

[gau-nernst]

Thank you for your reply. Look forward to that feature! I was looking through ViT implementations in HF transformers but sadly they don't support global average pooling.

[rwightman]

@gau-nernst testing this right now, fairly easy addition, yay

this is what it will look like in the config file https://huggingface.co/rwightman/test_vit_b16/blob/main/config.json

[rwightman]

@gau-nernst it's on the main branch now, any interest in giving it a try before I put out a new release?

Intended to be used like this...

model_args=dict(global_pool='avg', act_layer='silu', fc_norm=False)
mm = timm.create_model('vit_base_patch16_clip_224', **model_args)
mm.load_state_dict(...)
push_to_hf_hub(mm, 'rwightman/test_vit_b16', model_args=model_args)

[rwightman]

note that the input_size (img size + in chans) still gets handled via a special path through the pretrained_cfg, num_classes as well, model_args should cover all other args that you could pass to any given model class...

[gau-nernst]

It works perfectly for my HF-hosted model. Thank you for the feature!.