Importing VISSL models

[dannyrichy]

Hello,

I wanted to ask if there are any plans to bring the VISSL models into the TIMM framework. While they may not be SOTA anymore, having them available in TIMM would still be very helpful for many researchers, especially since VISSL is no longer officially maintained and managing the library has become a nightmare.

Thanks

[vincere-mori]

No official plans that I know of, but since VISSL backbones are just standard ResNets and ViTs you can load the weights manually into timm equivalents without too much pain.

For ResNet-based checkpoints (MoCo, SimCLR, DINO-ResNet), the VISSL checkpoint wraps the trunk under something like classy_state_dict.base_model.model.trunk._feature_blocks. You need to strip that prefix:

import timm
import torch

ckpt = torch.load("vissl_rn50.torch", map_location="cpu")
state = ckpt["classy_state_dict"]["base_model"]["model"]["trunk"]
state = {k.replace("_feature_blocks.", ""): v for k, v in state.items()}

model = timm.create_model("resnet50", pretrained=False, num_classes=0)
missing, unexpected = model.load_state_dict(state, strict=False)
print(missing)  # should only be head weights if num_classes=0

strict=False is fine here as long as the main conv and layer blocks load correctly. Print missing and verify it is just classifier head keys.

For ViT-based VISSL models, some DINO ViT weights are already in timm directly, e.g. vit_small_patch16_224.dino. Worth checking timm.list_models("*dino*") before doing manual loading.