Loading weights into timm models

[AlaaKhaddaj]

I have trained a bunch of Resnet50 and ViT-B-32, and I would like now to use these weights to create a new model using timm. Is there an easy way to achieve this?

Thanks.

[rwightman]

@AlaaKhaddaj if you used a timm- model, it's very easy, you just need to modify the state dict, and then re-save...

ckpt = torch.load('open_clip.pt')
sd_timm = {k.replace('module.visual.trunk.', ''): v for k, v in ckpt['state_dict'].items() if k.startswith('module.visual.trunk')}

The open-clip ViT model and ResNet model need remapping though, I don't have a remap for open_clip ResNet -> timm, but do have code for ViT. NOTE, the ViT model is different than the default timm vit, so needs to be one of the vit_xxx_patchxx_clip_xx models not the ones w/o the 'clip' in them, there is an extra norm near the beginning of the network.

https://github.com/rwightman/pytorch-image-models/blob/94a91598c30125f6d95b08c7f2edcafd72b7dfa6/timm/models/vision_transformer.py#L600-L632

[vturrisi]

Hey @rwightman, is there any way to convert the CoCa models to timm?

[rwightman]

@vturrisi I need to support the different pooling that was mentioned in the issue you filed re the seq len, timm vits only support standard token + proj or global avg pool + proj outputs… I do intend to do that, but need to finalize the design.

[vturrisi]

@rwightman I see. And regarding the other layers, do you have the mapping dict for them? I wanted to try this https://github.com/facebookresearch/ToMe on CoCa. They have a wrapper to replace the common timm blocks, so if I could map the CoCa's weights (without the final attention pooling) to timm, I could run ToMe and just copy-paste the current attention pooling.