Segmentation Pretrained Weights

[isaaccorley]

This PR addresses part of #1044 adds the ability to load pretrained weights from a backbone model e.g. ResNet into a semantic segmentation encoder. This works for the segmentation-models-pytorch Unet and DeepLabv3 implementations but not the FCN because we aren't using a backbone for that.

[adamjstewart]

This only supports ImageNet pretrained weights, which is pretty useless for us. We really want support for our weight enums like we have for regression/classification/byol. Could we manually load weights from state dict ourselves like we do with our timm-based models?

[adamjstewart]

Oh wait, I didn't read the full PR. Maybe this does cover all the features we want. You'll definitely want to update the docstring though.

[adamjstewart]

Also needs tests

[nilsleh]

Maybe also put a note in the docs that not all our pretrained models (i.e. the ViTs) will be compatible according to the smp docs or maybe there is a way around, not sure.

[isaaccorley]

The other thing I realized is that by default smp will use torchvision for resnet backbones e.g. resnet18. If you want to force using timm implementations you have to prefix the encoder names with tu- e.g. tu-resnet18.

[calebrob6]

We could just only allow timm pretrained backbones here

[isaaccorley]

This actually needs #1049 to be solved otherwise this only supports resnet backbones.

[isaaccorley]

We could just only allow timm pretrained backbones here

So do we want to prefix tu- to all the backbones or do we want the user to do it?

[adamjstewart]

That is a tough question. I would be fine with only supporting timm backbones for now, which would suggest prepending tu- for the user so that it matches other trainers. But there may come a time in which we want to support non-timm backbones because pretrained models are not always made with timm.

[calebrob6]

We could prefix "tu-" if it results in a valid timm backbone to give ourselves room to grow.

[calebrob6]

Ugh, also the horrific model checkpoint tests

[calebrob6]

We need this :)

[calebrob6]

@isaaccorley does this work with our custom ResNet weights?

[isaaccorley]

I haven't tried using the weights enums but in theory it should since the encoders are just resnet models.

[adamjstewart]

It should be the same loading code and docstring description as all of the other trainers.

[isaaccorley]

@calebrob6 I just made it work with the ResNet weights

from torchgeo.trainers import SemanticSegmentationTask
from torchgeo.models import ResNet50_Weights

model = SemanticSegmentationTask(
    model="unet",
    backbone="resnet50",
    weights=ResNet50_Weights.SENTINEL2_RGB_MOCO,
    in_channels=3,
    num_classes=2,
    loss="ce",
    ignore_index=0,
    learning_rate=3e-4,
    learning_rate_schedule_patience=5,
    freeze_backbone=False,
    freeze_decoder=False,
)

[calebrob6]

LGTM!