To use self-supervised ImageNet weights, start by finding a PyTorch Lightning checkpoint for the model you want to use. Then, update the LOAD_PATH variable in the selfsup_imagenet_init_pretraining.sh script with the path to your checkpoint.
We’re currently using weights from Lightly. Similar to our setup for supervised ImageNet initialization, we want to start with the ResNet weights while ignoring the classifier/projector weights (Even if we want to use them, usually their shapes do not match with ours + they use different networks etc.). The checkpoint_modifier.py script reads the configuration, removes or renames some of the keys from the chekpoint and saves the modified checkpoint with the _modified.ckpt suffix. So at the end, original checkpoint file will stay the same.
We load the weights using load_from_checkpoint(strict=False). Initially, it will try to match all the keys. You should see a warning indicating that some keys in the model cannot be found in the checkpoint—these are typically related to the projection head and classifier. In such cases, the missing weights are randomly initialized, which is expected. If another warning appears indicating that some keys in the checkpoint do not exist in the model, this means checkpoint_modifier.py is not working correctly and needs to be updated to remove these keys. While this mismatch is not critical, it would be better to clean up the extra keys from the checkpoint, and it is a MUST to properly rename the keys that the model needs to use.
checkpoint_modifier.py currently handles SimCLR and BYOL. We have to expand it for other models as other models usually use different projectors
LOAD_PATH
variable in theselfsup_imagenet_init_pretraining.sh
script with the path to your checkpoint.checkpoint_modifier.py
script reads the configuration, removes or renames some of the keys from the chekpoint and saves the modified checkpoint with the_modified.ckpt
suffix. So at the end, original checkpoint file will stay the same.load_from_checkpoint(strict=False)
. Initially, it will try to match all the keys. You should see a warning indicating that some keys in the model cannot be found in the checkpoint—these are typically related to the projection head and classifier. In such cases, the missing weights are randomly initialized, which is expected. If another warning appears indicating that some keys in the checkpoint do not exist in the model, this means checkpoint_modifier.py is not working correctly and needs to be updated to remove these keys. While this mismatch is not critical, it would be better to clean up the extra keys from the checkpoint, and it is a MUST to properly rename the keys that the model needs to use.checkpoint_modifier.py
currently handles SimCLR and BYOL. We have to expand it for other models as other models usually use different projectors