Repository: microsoft/torchgeo

Issue:
utils.load_state_dict() does not support ViT architecture
### Description

When using the new pretrained weights, the `load_state_dict()` function is looking for a "conv1" in the model to determine the expected input channels to the models. However, only ResNet Architecture weights have this. ViTs begin with "patch_embed". I think the trainer `test_classification.py` as well as the others are only testing with a ResNet18, so that is why the tests didn't catch it.

### Steps to reproduce

```python
from torchgeo.models import ViTSmall16_Weights
from torchgeo.trainers import ClassificationTask

task = ClassificationTask(
    model="vit_small_patch16_224",
    weights=ViTSmall16_Weights.SENTINEL2_ALL_DINO,
    num_classes=10,
    in_channels=13,
)
```

### Version

0.4.0


Discussion:
Good catch! We could either add this to our model tests or our trainer tests or both. 

Honestly, I would love to get rid of everything in `torchgeo/trainers/utils.py` and instead use the builtin functions in PyTorch. The only reason we override `load_state_dict` is for when num_classes or in_channels doesn't match the model weights, but we could choose to only support model weights with the correct in_channels and always remove the last layer so num_classes doesn't matter. Not sure if that would be realistic or not as a fix for this.
I think having the flexibility of `num_classes` and `in_channels` is a nice feature for running experiment as a user so I'd be in favor of keeping it. I am working on a PR for adding more ViT weights, so I could fix the issue there or do you want a separate PR?
Separate PR. New features have to wait until 0.5, but bug fixes can be backported to 0.4.1. 
I think we can fix this by checking the parent class of the model

For example:

```python
import timm

resnet = timm.create_model("resnet18")
vit = timm.create_model("vit_small_patch16_224")

type(resnet)
# <class 'timm.models.resnet.ResNet'>

type(vit)
# <class 'timm.models.vision_transformer.VisionTransformer'>
```

Thoughts?
We will eventually need to support far more models than just ResNet and ViT. Prob best to use `hasattr` and look for common names for the first layer. Hopefully there are only a few names for the first layer.
- resnet starts with `conv1`
- vit starts with `patch_embed`
- efficientnet starts with `conv_stem`

I don't have a solution yet but I don't think this will be the answer either. It may be easier to just force the user to input which backbone the state dict was originally from.
They do input the backbone, but they don't input which layer to look at to find in_channels. 
I was thinking of finding the in_channels by looking at the first named parameter in the model with `model.named_parameters()`, however, in the ViT case index 0 is a `cls_token` and not the patch_embed so cannot just taken the first index. Same goes for the state_dict. I suppose one could generate a dict that maps all backbone names in timm to the named parameter where to look for the in_channels... not very elegant.
That's actually a really interesting idea! I would do that as the fallback and then only special case the few architectures like ViT where that doesn't work.
ViT `patch_embed` actually starts with a conv layer. So if you do `list(list(model.children())[0].children())[0]` you can access it.

If you recursively search the first named children until it's a base module with no children you can actually get the key e.g. and this works for resnet, vit, efficientnet, etc. Haven't done a thorough search though.

```python
import timm

def get_input_layer(backbone):
    model = timm.create_model(backbone)

    keys = []
    children = list(model.named_children())
    while children != []:
        name, module = children[0]
        keys.append(name)
        children = list(module.named_children())
    
    key = ".".join(keys)
    return key, module

get_input_layer("resnet18")
# ('conv1',
# Conv2d(3, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False))

get_input_layer("vit_small_patch16_224")
# ('patch_embed.proj', Conv2d(3, 384, kernel_size=(16, 16), stride=(16, 16)))

get_input_layer("efficientnet_b0")
# ('conv_stem',
# Conv2d(3, 32, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False))

```
```
import timm
import pandas as pd


def get_input_layer(backbone):
    model = timm.create_model(backbone)

    keys = []
    children = list(model.named_children())
    while children != []:
        name, module = children[0]
        keys.append(name)
        children = list(module.named_children())

    key = ".".join(keys)
    return key, module


all_timm_model_names = timm.list_models()
result = {}
for model_name in tqdm(all_timm_model_names):
    key, module = get_input_layer(model_name)
    try:
        num_in_channels = module.in_channels
        result[model_name] = num_in_channels
    except:
        result[model_name] = -1

df = (
    pd.DataFrame.from_dict(result, orient="index")
    .reset_index()
    .rename(columns={0: "num_in_chans", "index": "model_name"})
)

print(df[df["num_in_chans"]==-1])
```

If I loop through all timm models with your code, all models have `num_in_channels=3`, except for:
```
               model_name  num_in_chans
532             tresnet_l            -1
533         tresnet_l_448            -1
534             tresnet_m            -1
535         tresnet_m_448            -1
536  tresnet_m_miil_in21k            -1
537            tresnet_xl            -1
538        tresnet_xl_448            -1
```
Dang you beat me. I'm just finishing looping through them in the next minute or so
If all models default to 3 channels then we can just hardcode it to 3
I believe all models are trained for imagenet or variants of imagenet so it's unlikely that it's not 3 channels.
But I think you still need to obtain the first parameter name to index the state_dict to get `expected_in_channels` from the weights you try to load and for that you need Isaac's solution.
Some analysis on the type of the first layer. Majority is `Conv2d`. Looks like the only one that is not a subclass of `Conv2d` is `SpaceToDepthModule` which doesn't depend on input channels.

I think it's fair to assume that we can use the above solution to get the state_dict key and input channels.

```
Counter({
         'Conv2d': 825,
         'ScaledStdConv2dSame': 7,
         'ScaledStdConv2d': 28,
         'Conv2dSame': 66,
         'StdConv2d': 15,
         'SpaceToDepthModule': 8,
         'StdConv2dSame': 15
})
```

Resolving patch:
diff --git a/torchgeo/trainers/byol.py b/torchgeo/trainers/byol.py
--- a/torchgeo/trainers/byol.py
+++ b/torchgeo/trainers/byol.py
@@ -313,7 +313,7 @@ def config_task(self) -> None:
                 state_dict = get_weight(weights).get_state_dict(progress=True)
             backbone = utils.load_state_dict(backbone, state_dict)
 
-        self.model = BYOL(backbone, in_channels=in_channels, image_size=(256, 256))
+        self.model = BYOL(backbone, in_channels=in_channels, image_size=(224, 224))
 
     def __init__(self, **kwargs: Any) -> None:
         """Initialize a LightningModule for pre-training a model with BYOL.
diff --git a/torchgeo/trainers/utils.py b/torchgeo/trainers/utils.py
--- a/torchgeo/trainers/utils.py
+++ b/torchgeo/trainers/utils.py
@@ -54,6 +54,23 @@ def extract_backbone(path: str) -> Tuple[str, "OrderedDict[str, Tensor]"]:
     return name, state_dict
 
 
+def _get_input_layer_name_and_module(model: Module) -> Tuple[str, Module]:
+    """Retrieve the input layer name and module from a timm model.
+
+    Args:
+        model: timm model
+    """
+    keys = []
+    children = list(model.named_children())
+    while children != []:
+        name, module = children[0]
+        keys.append(name)
+        children = list(module.named_children())
+
+    key = ".".join(keys)
+    return key, module
+
+
 def load_state_dict(model: Module, state_dict: "OrderedDict[str, Tensor]") -> Module:
     """Load pretrained resnet weights to a model.
 
@@ -68,27 +85,32 @@ def load_state_dict(model: Module, state_dict: "OrderedDict[str, Tensor]") -> Mo
         If input channels in model != pretrained model input channels
         If num output classes in model != pretrained model num classes
     """
-    in_channels = cast(nn.Module, model.conv1).in_channels
-    expected_in_channels = state_dict["conv1.weight"].shape[1]
+    input_module_key, input_module = _get_input_layer_name_and_module(model)
+    in_channels = input_module.in_channels
+    expected_in_channels = state_dict[input_module_key + ".weight"].shape[1]
 
-    num_classes = cast(nn.Module, model.fc).out_features
+    output_module_key, output_module = list(model.named_children())[-1]
+    num_classes = output_module.out_features
     expected_num_classes = None
-    if "fc.weight" in state_dict:
-        expected_num_classes = state_dict["fc.weight"].shape[0]
+    if output_module_key + ".weight" in state_dict:
+        expected_num_classes = state_dict[output_module_key + ".weight"].shape[0]
 
     if in_channels != expected_in_channels:
         warnings.warn(
             f"input channels {in_channels} != input channels in pretrained"
             f" model {expected_in_channels}. Overriding with new input channels"
         )
-        del state_dict["conv1.weight"]
+        del state_dict[input_module_key + ".weight"]
 
     if expected_num_classes and num_classes != expected_num_classes:
         warnings.warn(
             f"num classes {num_classes} != num classes in pretrained model"
             f" {expected_num_classes}. Overriding with new num classes"
         )
-        del state_dict["fc.weight"], state_dict["fc.bias"]
+        del (
+            state_dict[output_module_key + ".weight"],
+            state_dict[output_module_key + ".bias"],
+        )
 
     model.load_state_dict(state_dict, strict=False)
     return model


Test patch:
diff --git a/tests/trainers/test_byol.py b/tests/trainers/test_byol.py
--- a/tests/trainers/test_byol.py
+++ b/tests/trainers/test_byol.py
@@ -10,6 +10,7 @@
 import torch
 import torch.nn as nn
 import torchvision
+from _pytest.fixtures import SubRequest
 from _pytest.monkeypatch import MonkeyPatch
 from omegaconf import OmegaConf
 from pytorch_lightning import LightningDataModule, Trainer
@@ -18,7 +19,7 @@
 
 from torchgeo.datamodules import ChesapeakeCVPRDataModule, MisconfigurationException
 from torchgeo.datasets import ChesapeakeCVPR
-from torchgeo.models import ResNet18_Weights
+from torchgeo.models import get_model_weights, list_models
 from torchgeo.samplers import GridGeoSampler
 from torchgeo.trainers import BYOLTask
 from torchgeo.trainers.byol import BYOL, SimCLRAugmentation
@@ -97,13 +98,30 @@ def test_trainer(
 
     @pytest.fixture
     def model_kwargs(self) -> Dict[str, Any]:
-        return {"backbone": "resnet18", "weights": None, "in_channels": 3}
+        return {
+            "backbone": "resnet18",
+            "in_channels": 13,
+            "loss": "ce",
+            "num_classes": 10,
+            "weights": None,
+        }
+
+    @pytest.fixture(
+        params=[
+            weights for model in list_models() for weights in get_model_weights(model)
+        ]
+    )
+    def weights(self, request: SubRequest) -> WeightsEnum:
+        return request.param
 
     @pytest.fixture
-    def mocked_weights(self, tmp_path: Path, monkeypatch: MonkeyPatch) -> WeightsEnum:
-        weights = ResNet18_Weights.SENTINEL2_RGB_MOCO
+    def mocked_weights(
+        self, tmp_path: Path, monkeypatch: MonkeyPatch, weights: WeightsEnum
+    ) -> WeightsEnum:
         path = tmp_path / f"{weights}.pth"
-        model = timm.create_model("resnet18", in_chans=weights.meta["in_chans"])
+        model = timm.create_model(
+            weights.meta["model"], in_chans=weights.meta["in_chans"]
+        )
         torch.save(model.state_dict(), path)
         monkeypatch.setattr(weights, "url", str(path))
         monkeypatch.setattr(torchvision.models._api, "load_state_dict_from_url", load)
@@ -111,20 +129,43 @@ def mocked_weights(self, tmp_path: Path, monkeypatch: MonkeyPatch) -> WeightsEnu
 
     def test_weight_file(self, model_kwargs: Dict[str, Any], checkpoint: str) -> None:
         model_kwargs["weights"] = checkpoint
-        BYOLTask(**model_kwargs)
+        with pytest.warns(UserWarning):
+            BYOLTask(**model_kwargs)
 
     def test_weight_enum(
         self, model_kwargs: Dict[str, Any], mocked_weights: WeightsEnum
     ) -> None:
+        model_kwargs["backbone"] = mocked_weights.meta["model"]
+        model_kwargs["in_channels"] = mocked_weights.meta["in_chans"]
         model_kwargs["weights"] = mocked_weights
         BYOLTask(**model_kwargs)
 
     def test_weight_str(
         self, model_kwargs: Dict[str, Any], mocked_weights: WeightsEnum
     ) -> None:
+        model_kwargs["backbone"] = mocked_weights.meta["model"]
+        model_kwargs["in_channels"] = mocked_weights.meta["in_chans"]
         model_kwargs["weights"] = str(mocked_weights)
         BYOLTask(**model_kwargs)
 
+    @pytest.mark.slow
+    def test_weight_enum_download(
+        self, model_kwargs: Dict[str, Any], weights: WeightsEnum
+    ) -> None:
+        model_kwargs["backbone"] = weights.meta["model"]
+        model_kwargs["in_channels"] = weights.meta["in_chans"]
+        model_kwargs["weights"] = weights
+        BYOLTask(**model_kwargs)
+
+    @pytest.mark.slow
+    def test_weight_str_download(
+        self, model_kwargs: Dict[str, Any], weights: WeightsEnum
+    ) -> None:
+        model_kwargs["backbone"] = weights.meta["model"]
+        model_kwargs["in_channels"] = weights.meta["in_chans"]
+        model_kwargs["weights"] = str(weights)
+        BYOLTask(**model_kwargs)
+
     def test_predict(self, model_kwargs: Dict[Any, Any], fast_dev_run: bool) -> None:
         datamodule = PredictBYOLDataModule(
             root="tests/data/chesapeake/cvpr",
diff --git a/tests/trainers/test_classification.py b/tests/trainers/test_classification.py
--- a/tests/trainers/test_classification.py
+++ b/tests/trainers/test_classification.py
@@ -9,6 +9,7 @@
 import timm
 import torch
 import torchvision
+from _pytest.fixtures import SubRequest
 from _pytest.monkeypatch import MonkeyPatch
 from omegaconf import OmegaConf
 from pytorch_lightning import LightningDataModule, Trainer
@@ -24,7 +25,7 @@
     UCMercedDataModule,
 )
 from torchgeo.datasets import BigEarthNet, EuroSAT
-from torchgeo.models import ResNet18_Weights
+from torchgeo.models import get_model_weights, list_models
 from torchgeo.trainers import ClassificationTask, MultiLabelClassificationTask
 
 from .test_utils import ClassificationTestModel
@@ -110,11 +111,22 @@ def model_kwargs(self) -> Dict[str, Any]:
             "weights": None,
         }
 
+    @pytest.fixture(
+        params=[
+            weights for model in list_models() for weights in get_model_weights(model)
+        ]
+    )
+    def weights(self, request: SubRequest) -> WeightsEnum:
+        return request.param
+
     @pytest.fixture
-    def mocked_weights(self, tmp_path: Path, monkeypatch: MonkeyPatch) -> WeightsEnum:
-        weights = ResNet18_Weights.SENTINEL2_ALL_MOCO
+    def mocked_weights(
+        self, tmp_path: Path, monkeypatch: MonkeyPatch, weights: WeightsEnum
+    ) -> WeightsEnum:
         path = tmp_path / f"{weights}.pth"
-        model = timm.create_model("resnet18", in_chans=weights.meta["in_chans"])
+        model = timm.create_model(
+            weights.meta["model"], in_chans=weights.meta["in_chans"]
+        )
         torch.save(model.state_dict(), path)
         monkeypatch.setattr(weights, "url", str(path))
         monkeypatch.setattr(torchvision.models._api, "load_state_dict_from_url", load)
@@ -128,6 +140,8 @@ def test_weight_file(self, model_kwargs: Dict[str, Any], checkpoint: str) -> Non
     def test_weight_enum(
         self, model_kwargs: Dict[str, Any], mocked_weights: WeightsEnum
     ) -> None:
+        model_kwargs["model"] = mocked_weights.meta["model"]
+        model_kwargs["in_channels"] = mocked_weights.meta["in_chans"]
         model_kwargs["weights"] = mocked_weights
         with pytest.warns(UserWarning):
             ClassificationTask(**model_kwargs)
@@ -135,10 +149,30 @@ def test_weight_enum(
     def test_weight_str(
         self, model_kwargs: Dict[str, Any], mocked_weights: WeightsEnum
     ) -> None:
+        model_kwargs["model"] = mocked_weights.meta["model"]
+        model_kwargs["in_channels"] = mocked_weights.meta["in_chans"]
         model_kwargs["weights"] = str(mocked_weights)
         with pytest.warns(UserWarning):
             ClassificationTask(**model_kwargs)
 
+    @pytest.mark.slow
+    def test_weight_enum_download(
+        self, model_kwargs: Dict[str, Any], weights: WeightsEnum
+    ) -> None:
+        model_kwargs["model"] = weights.meta["model"]
+        model_kwargs["in_channels"] = weights.meta["in_chans"]
+        model_kwargs["weights"] = weights
+        ClassificationTask(**model_kwargs)
+
+    @pytest.mark.slow
+    def test_weight_str_download(
+        self, model_kwargs: Dict[str, Any], weights: WeightsEnum
+    ) -> None:
+        model_kwargs["model"] = weights.meta["model"]
+        model_kwargs["in_channels"] = weights.meta["in_chans"]
+        model_kwargs["weights"] = str(weights)
+        ClassificationTask(**model_kwargs)
+
     def test_invalid_loss(self, model_kwargs: Dict[str, Any]) -> None:
         model_kwargs["loss"] = "invalid_loss"
         match = "Loss type 'invalid_loss' is not valid."
diff --git a/tests/trainers/test_regression.py b/tests/trainers/test_regression.py
--- a/tests/trainers/test_regression.py
+++ b/tests/trainers/test_regression.py
@@ -9,6 +9,7 @@
 import timm
 import torch
 import torchvision
+from _pytest.fixtures import SubRequest
 from _pytest.monkeypatch import MonkeyPatch
 from omegaconf import OmegaConf
 from pytorch_lightning import LightningDataModule, Trainer
@@ -20,7 +21,7 @@
     TropicalCycloneDataModule,
 )
 from torchgeo.datasets import TropicalCyclone
-from torchgeo.models import ResNet18_Weights
+from torchgeo.models import get_model_weights, list_models
 from torchgeo.trainers import RegressionTask
 
 from .test_utils import RegressionTestModel
@@ -86,11 +87,22 @@ def model_kwargs(self) -> Dict[str, Any]:
             "in_channels": 3,
         }
 
+    @pytest.fixture(
+        params=[
+            weights for model in list_models() for weights in get_model_weights(model)
+        ]
+    )
+    def weights(self, request: SubRequest) -> WeightsEnum:
+        return request.param
+
     @pytest.fixture
-    def mocked_weights(self, tmp_path: Path, monkeypatch: MonkeyPatch) -> WeightsEnum:
-        weights = ResNet18_Weights.SENTINEL2_RGB_MOCO
+    def mocked_weights(
+        self, tmp_path: Path, monkeypatch: MonkeyPatch, weights: WeightsEnum
+    ) -> WeightsEnum:
         path = tmp_path / f"{weights}.pth"
-        model = timm.create_model("resnet18", in_chans=weights.meta["in_chans"])
+        model = timm.create_model(
+            weights.meta["model"], in_chans=weights.meta["in_chans"]
+        )
         torch.save(model.state_dict(), path)
         monkeypatch.setattr(weights, "url", str(path))
         monkeypatch.setattr(torchvision.models._api, "load_state_dict_from_url", load)
@@ -104,6 +116,8 @@ def test_weight_file(self, model_kwargs: Dict[str, Any], checkpoint: str) -> Non
     def test_weight_enum(
         self, model_kwargs: Dict[str, Any], mocked_weights: WeightsEnum
     ) -> None:
+        model_kwargs["model"] = mocked_weights.meta["model"]
+        model_kwargs["in_channels"] = mocked_weights.meta["in_chans"]
         model_kwargs["weights"] = mocked_weights
         with pytest.warns(UserWarning):
             RegressionTask(**model_kwargs)
@@ -111,10 +125,30 @@ def test_weight_enum(
     def test_weight_str(
         self, model_kwargs: Dict[str, Any], mocked_weights: WeightsEnum
     ) -> None:
+        model_kwargs["model"] = mocked_weights.meta["model"]
+        model_kwargs["in_channels"] = mocked_weights.meta["in_chans"]
         model_kwargs["weights"] = str(mocked_weights)
         with pytest.warns(UserWarning):
             RegressionTask(**model_kwargs)
 
+    @pytest.mark.slow
+    def test_weight_enum_download(
+        self, model_kwargs: Dict[str, Any], weights: WeightsEnum
+    ) -> None:
+        model_kwargs["model"] = weights.meta["model"]
+        model_kwargs["in_channels"] = weights.meta["in_chans"]
+        model_kwargs["weights"] = weights
+        RegressionTask(**model_kwargs)
+
+    @pytest.mark.slow
+    def test_weight_str_download(
+        self, model_kwargs: Dict[str, Any], weights: WeightsEnum
+    ) -> None:
+        model_kwargs["model"] = weights.meta["model"]
+        model_kwargs["in_channels"] = weights.meta["in_chans"]
+        model_kwargs["weights"] = str(weights)
+        RegressionTask(**model_kwargs)
+
     def test_no_rgb(
         self, monkeypatch: MonkeyPatch, model_kwargs: Dict[Any, Any], fast_dev_run: bool
     ) -> None:
diff --git a/tests/trainers/test_utils.py b/tests/trainers/test_utils.py
--- a/tests/trainers/test_utils.py
+++ b/tests/trainers/test_utils.py
@@ -6,11 +6,13 @@
 from typing import Any, cast
 
 import pytest
+import timm
 import torch
 import torch.nn as nn
 from torch.nn.modules import Module
 
 from torchgeo.trainers.utils import (
+    _get_input_layer_name_and_module,
     extract_backbone,
     load_state_dict,
     reinit_initial_conv_layer,
@@ -115,3 +117,10 @@ def test_reinit_initial_conv_layer() -> None:
     assert in_channels == 4
     assert k1 == 3 and k2 == 3
     assert new_conv_layer.stride[0] == 2
+
+
+def test_get_input_layer_name_and_module() -> None:
+    key, module = _get_input_layer_name_and_module(timm.create_model("resnet18"))
+    assert key == "conv1"
+    assert isinstance(module, nn.Conv2d)
+    assert module.in_channels == 3