Issue in loading weights of ResNet

[LucaBarco]

Good morning,
I'm facing the following issue trying to finetune SSL4EO ResNet18 model.

I'm executing this code

    backbone = "terratorch_ssl4eos12_resnet50_sentinel2_rgb_moco",
    backbone_pretrained = True,
    backbone_model_bands = [
         "R", "G", "B"
    ],
    decoder = "FCNDecoder",
    decoder_channels = 512,
    head_dropout = 0.1,
    num_classes = 2)

task = SemanticSegmentationTask(model_args,
                                model_factory="EncoderDecoderFactory",                               
                                freeze_backbone=True, freeze_decoder=False, optimizer="AdamW", lr=1e-3, class_weights=[0.53427446, 7.7940613 ],
    optimizer_hparams={"weight_decay": 0.05})

I receive an error during the creation of SemanticSegmentationTask object due to a name mismatch between pretrained ResNet Layers names (that don't contain "0", i.e. B1) and the lookup table present in line 62 in terratoch/models/backbones/torchgeo_resnet.py that contains "0" (i.e. B01):

look_up_table = {
    "B01": "COASTAL_AEROSOL",
    "B02": "BLUE",
    "B03": "GREEN",
    "B04": "RED",
    "B05": "RED_EDGE_1",
    "B06": "RED_EDGE_2",
    "B07": "RED_EDGE_3",
    "B08": "NIR_BROAD",
    "B8A": "NIR_NARROW",
    "B09": "WATER_VAPOR",
    "B10": "CIRRUS",
    "B11": "SWIR_1",
    "B12": "SWIR_2",
    "VV": "VV",
    "VH": "VH",
    "R": "RED",
    "G": "GREEN",
    "B": "BLUE"
}

the
'`

Screenshots or log output

Log Output

`---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
[/tmp/ipython-input-1465269867.py](https://localhost:8080/#) in ()
     28 )
     29 
---> 30 task = SemanticSegmentationTask(model_args,
     31                                 model_factory="EncoderDecoderFactory",
     32                                 freeze_backbone=True, freeze_decoder=False, optimizer="AdamW", lr=1e-3, class_weights=[0.53427446, 7.7940613 ])
10 frames

/usr/local/lib/python3.12/dist-packages/terratorch/models/backbones/torchgeo_resnet.py in get_pretrained_bands(model_bands)

96 def get_pretrained_bands(model_bands):

97

---> 98     model_bands = [look_up_table[x.split('.')[-1]] for x in model_bands]

99

100     return model_bands
KeyError: 'B1

Expected behavior
The SemanticSegmentationTask object should be instantiated without this error.

Deployment information
Describe what you've deployed and how:

• TerraTorch version: 1.1
• Installation source: pip

[Joao-L-S-Almeida]

@LucaBarco It looks like some of the components form this module haven't being properly tested in pretrained model. I adapted your script as:

from terratorch.tasks import SemanticSegmentationTask

model_args = {
    "backbone": "terratorch_ssl4eos12_resnet50_sentinel2_rgb_moco",
    "backbone_pretrained": True, #False, 
    "backbone_model_bands": [
        "B02", "B03", "B04"
    ],
    "decoder":  "FCNDecoder",
    "decoder_channels": 512,
    "head_dropout": 0.1,
    "num_classes": 2
}

task = SemanticSegmentationTask(model_args,
                                model_factory="EncoderDecoderFactory",                               
                                freeze_backbone=True, freeze_decoder=False, optimizer="AdamW", lr=1e-3, class_weights=[0.53427446, 7.7940613 ],
    optimizer_hparams={"weight_decay": 0.05})

print(task)

And I had do add new entries in look_up_table because in the metadata for this model the bands are named as B<number>. It doesn't seem it has changed in torchgeo, it looks like it was never tested and the entries are really missing.
Any comment @romeokienzler @paolofraccaro @blumenstiel ?