Fix load_state_dict for all timm models

[nilsleh]

This PR closes #1049 , by implementing Isaac's solution.

[nilsleh]

I looked into the seco weights again. Since they are originally saved as part of a pytorch-lightning module, the keys have different names then the default timm keys. Looking at the seco code there are a "q" and a "k" network and the "q" network is used as a pretrained backbone for downstream tasks. The encoder_q.0.weight is the only weight that has the same shape as the default timm model and there is also no bias at this stage, so I think encoder_q.0.weight is the conv1.weight we look for. I can rename it and reupload to huggingface and then I think we should be good. Or do you have a suggestion?

import torch
import timm

timm_model = timm.create_model("resnet18")
checkpoint = torch.load("path/to/seco_resnet18_1m.ckpt", map_location="cpu")

state_dict = checkpoint["state_dict"]
assert timm_model.conv1.weight.shape == state_dict["encoder_q.0.weight"].shape

for key, val in checkpoint["state_dict"].items():
    if key not in ["encoder_q.0.weight", "encoder_k.0.weight"]:
        if state_dict["encoder_q.0.weight"].shape == val.shape:
            print(key)

state_dict["conv1.weight"] = state_dict.pop("encoder_q.0.weight")
# save this state dict with just the encoder_q weights and renamed 0.weight and upload to huggingface

[adamjstewart]

Your guess is as good as mine. Do the authors have any code for loading the pretrained model like your group does? If not, then I think your analysis makes sense. If you really want to be sure, you can try to train a model with those pretrained weights and make sure if converges quickly.

[calebrob6]

You could use something like this -- https://gist.github.com/calebrob6/44f2e42017e2d192e837f0a1cd526c50 -- to make sure that linear-probing on a downstream dataset with the model achieves good performance. This is the notebook that I used to verify one of the SSL4EO weights I think.

[nilsleh]

Yes, so the encoder_q I got from this code where they load a backbone from their pretrained model for a downstream task.

[nilsleh]

I did the above described extraction of the weights and tried @calebrob6 script with the extracted seco resnet18 weights. As they are pretrained on RGB I only use the Eurosat RGB bands. Here are the scores I get:

Edit:
Correction, I forgot to call model.eval()...

This is with preprocessing step of just dividing image values by 1000:

• seco resnet18 weights 0.7625
• timm resnet imagenet 0.79
• timm resnet random init 0.46

This is with preprocessing step of using the provided normalization stats for bands

• seco resnet18 weights 0.75
• timm resnet imagenet 0.89
• timm resnet random init 0.65

[nilsleh]

I will try to investigate again tomorrow. This is what I am using as a script (takes about 3 minutes to run on cpu locally). And for some reason one needs pytorch-lightning==1.1.8 in order to do load the original checkpoint file.

[nilsleh]

I downloaded the bigearthnet dataset to try the linear probing script on that since that is the dataset they also report in their paper. For bigearthnet with 10,000 samples (not the full dataset), I get the following scores:

• seco resnet18 weights: 0.4983
• timm imagenet: 0.4512
• timm random: 0.442

[adamjstewart]

Interesting evidence that the weights may not be very transferable...

[adamjstewart]

But this should be sufficient proof that your approach to extracting the first layer of weights is correct and we can move forward with this PR.

[nilsleh]

In the paper, they report quiet a significant improvement when using Seco in linear probing, so I must be doing something wrong. I can also contact the authors to sort it out.

[adamjstewart]

That's prob for MSI, not RGB

[nilsleh]

I think everything is only RGB, at least that is how I interpret it when they state "Although the collected dataset contains up to 12 spectral bands, in this work we focus on the RGB channels since it is a more general modality."

[nilsleh]

For the moment I updated the seco weights on huggingface, and the loading works for all weights now.

[adamjstewart]

I'm still seeing the same issue:

$ pytest -m slow tests/trainers/
...
FAILED tests/trainers/test_byol.py::TestBYOLTask::test_weight_enum_download[ResNet18_Weights.SENTINEL2_RGB_SECO] - KeyError: 'conv1.weight'
FAILED tests/trainers/test_byol.py::TestBYOLTask::test_weight_enum_download[ResNet50_Weights.SENTINEL2_RGB_SECO] - KeyError: 'conv1.weight'
FAILED tests/trainers/test_byol.py::TestBYOLTask::test_weight_str_download[ResNet18_Weights.SENTINEL2_RGB_SECO] - KeyError: 'conv1.weight'
FAILED tests/trainers/test_byol.py::TestBYOLTask::test_weight_str_download[ResNet50_Weights.SENTINEL2_RGB_SECO] - KeyError: 'conv1.weight'
FAILED tests/trainers/test_classification.py::TestClassificationTask::test_weight_enum_download[ResNet18_Weights.SENTINEL2_RGB_SECO] - KeyError: 'conv1.weight'
FAILED tests/trainers/test_classification.py::TestClassificationTask::test_weight_enum_download[ResNet50_Weights.SENTINEL2_RGB_SECO] - KeyError: 'conv1.weight'
FAILED tests/trainers/test_classification.py::TestClassificationTask::test_weight_str_download[ResNet18_Weights.SENTINEL2_RGB_SECO] - KeyError: 'conv1.weight'
FAILED tests/trainers/test_classification.py::TestClassificationTask::test_weight_str_download[ResNet50_Weights.SENTINEL2_RGB_SECO] - KeyError: 'conv1.weight'
FAILED tests/trainers/test_regression.py::TestRegressionTask::test_weight_enum_download[ResNet18_Weights.SENTINEL2_RGB_SECO] - KeyError: 'conv1.weight'
FAILED tests/trainers/test_regression.py::TestRegressionTask::test_weight_enum_download[ResNet50_Weights.SENTINEL2_RGB_SECO] - KeyError: 'conv1.weight'
FAILED tests/trainers/test_regression.py::TestRegressionTask::test_weight_str_download[ResNet18_Weights.SENTINEL2_RGB_SECO] - KeyError: 'conv1.weight'
FAILED tests/trainers/test_regression.py::TestRegressionTask::test_weight_str_download[ResNet50_Weights.SENTINEL2_RGB_SECO] - KeyError: 'conv1.weight'

[nilsleh]

Mhm I am not getting those errors. Maybe, the old weights are still cached in your torch/hub? I had to delete those so it would reload the new ones after I uploaded them to huggingface.

[adamjstewart]

Oh mine are prob cached, let me delete

[adamjstewart]

Yep, works now. Thanks!