Load Segmentation Trainer Weights

[nilsleh]

This PR changes the loading of a checkpoint in the segmentation model. Given that I have trained a model with a torchgeo trainer, I might want to do with the weights argument:

• load the entire model (encoder and decoder)
• load just a backbone
• use only the model backbone from a trained model in a different task or at least make it available for other tasks

Not sure which of all should be supported by default, or whether there should be suggestions on how to do each of these (or other things).

[adamjstewart]

Need to think about this one. I wonder if there isn't an easier way to tell Lightning to save the encoder in a separate file. Then it's trainer-dependent instead of having to if-statement every possible naming scheme in one function.

[isaaccorley]

I'm not sure I understand. You are trying to load a model trained using which trainer into the segmentation trainer? Can you give the error message?

[isaaccorley]

If you need to load a checkpoint in the same trainer you can just use pytorch Lightning's loading method.

SemanticSegmentationTask.load_from_checkpoint(ckpt_path)

[adamjstewart]

I believe the checkpoint comes from MoCo and/or SimCLR and looks very different than the checkpoint expected by SemanticSegmentationTask.

[isaaccorley]

I don't think so. MoCo and SimCLR tasks named the backbone "backbone" not "encoder.

[isaaccorley]

Ah I see what's happening. Right now we don't support loading an entire pretrained segmentation checkpoint using the weights option. We only support loading the encoder part of the weights, not the decoder. So the solution would be to just load the checkpoint directly using pytorch lighting's checkpoint option.

Edit: so in this case, the PR changes are the correct solution.

[adamjstewart]

Is this still needed or is this superseded by #1403?

[isaaccorley]

I believe this is separate. This has to do with loading an entire UNet checkpoint using the weights argument. Right now the weights argument only loads the weights into the encoder (not the decoder).