Naming Scheme multi-weights pretrained models

[nilsleh]

This discussion concerns #762 and support for multi-weight pretrained models following the recently introduced torchvision multi-weight api. Specifically, this discussion should focus on a naming scheme for already available or future pretrained weights. I suppose for TorchGeo a name could include the dataset on which model was pretrained on and band specifics (only rgb or multispectral), as well as the satellite source.

[adamjstewart]

Attributes

The first thing we should think about is what attributes we consider to be important. This dictates what goes into the naming scheme.

1. model: we obviously have to include the model name somewhere
2. sensor: the imagery that the model is trained on controls everything from spectral bands to image resolution
3. bands: we'll probably want the option to provide models trained on all bands or a subset of bands
4. dataset?: do we want to be able to provide models trained on different datasets?
5. training method?: do we want to be able to provide models trained via SL or SSL? Different SSL approaches?
6. task?: does a model pretrained on land cover mapping work better for land cover mapping than a model trained on something else?
7. version? do we want to be able to provide updated versions of models in the future?

Anything else we could potentially consider important?

1 and 2 are obviously required, you can't load weights if they are from a different model, and you can't transfer learn from one sensor to the next (although maybe you could with RGB-only bands...). @calebrob6 wants to be able to provide models trained on all spectral bands or a subset of those bands (RGB-only, RGB+NIR, etc.). Torchvision uses 4 in their naming scheme. Based on https://github.com/zhu-xlab/SSL4EO-S12#pre-trained-models, it seems that 5 can affect model performance (i.e., you may want to choose a different SSL method depending on which dataset you are using). 6 is unclear to me and could be subsumed into 4. Torchvision uses 7 in their naming scheme to provide updated model weights without breaking backwards compatibility.

In addition to the question of what are all possible things we could care about, we also need to think about how many things we should care about. I want TorchGeo to be easy to use for a remote sensing expert who knows very little about ML. Users shouldn't have to research various SSL methods in order to decide which weights to use. I want to try to keep things as simple as possible so someone can say "I'm working with Sentinel-2 data, give me your recommended weights". This could come in the form of a table that provides model performance on various 1–7 combinations, or optional parameters that allow you to only specify the things you care about, or shortened aliases marked as default/recommended/latest.

[nilsleh]

I think either you have to spell out those given points in the name itself to make sure you have enough flexibility to name pretrained models that could come from various backgrounds or create an identifier that links to a documentation where all the details about the pretraining regime are listed.

[adamjstewart]

Model API

The next thing we should consider is the API with which we present our weights. Here are some options.

Current API

Our current API mimics the old torchvision design, with additional parameters for sensor/bands:

model = resnet50(sensor="sentinel2", bands="all", pretrained=True)

Pros:

• Allows each parameter to be optional (have a recommended default)
• Clear naming scheme (don't have to guess context of what "all" means in the middle of an enum)
• Relatively easy to add new parameters at a later date without changing the naming scheme
• Allows you to control # input channels even if you aren't using pretrained models

Cons:

• Not all combinations of all parameters will be supported (harder to document which are supported)
• Unclear exactly how models were trained (no additional metadata)
• Reproducibility is a problem unless users specify all parameters, even the optional ones or newly added ones

Enums

We could adopt an API that mimics the new torchvision design with enums:

model = resnet50(weights=ResNet50_Weights.SENTINEL2_ALL_DEFAULT)

Pros:

• Reproducibility, as long as users don't use DEFAULT
• Easier to document what weights are available (all valid enum values)
• No need to validate input (if the enum doesn't exist, you obviously can't use it as input)

Cons:

• Not easy to explain what each parts of the naming scheme means (what does ALL mean?)
• Not easy to add a new attribute later (do we rename existing or give up on consistency?)
• Duplicate model name (once in function, again in weight)
• No control over input channels unless you decide to use pretrained models
• All attributes are required (you could add a DEFAULT, but only at the end)

Other

Do any other libraries (timm, smp) allow you to select which model weights to use? How do they do it? Is there some kind of hybrid of the above two that would work better? Is there a completely different technique we could use, like some kind of nested enum with defaults at each level?

I kind of like the idea of having a single model function that takes the model ("resnet50") as input. This would be much easier to use in our trainers, since you can only easily store strings and floats in YAML.

[adamjstewart]

A counter argument against our current API is that it becomes hard to ensure that the default values will be valid as soon as you change one of the arguments. E.g., if I switch from sentinel2 to landsat8, there may not be the same set of models trained or tasks trained on.

[nilsleh]

You could, but as soon as someone cares about one of the other attributes, they're now forced to care about all of the other attributes.

Yes, certainly, but maybe that strikes a balance between a quick accessible baseline, and diving into specifics to use models. Yes, it is a long name to type but at the same time it will probably only be a few people and the vast majority just uses one of the above defaults, which seems like a reasonable naming scheme to me.

[nilsleh]

A counter argument against our current API is that it becomes hard to ensure that the default values will be valid as soon as you change one of the arguments. E.g., if I switch from sentinel2 to landsat8, there may not be the same set of models trained or tasks trained on.

Maybe then having just weights and num_input_channels as arguments is also a reasonable approach, as it follows the timm and smp setup. Then it is still clear that only the weights options are valid but you can have various input channels for your use case. The downside would be that you have to make decisions about how to extend channels or use fewer channels than the loaded weights have.

[nilsleh]

Coming back to this conversation, because I want to try an implementation for this. Could one idea be to include 1, 2, 3 as a minimum into the naming scheme followed by a "unique identifier" that can be quiet flexible and then have detailed documentation, about all other parameters, ssl methods, datasets and any additional information?

[adamjstewart]

Does this help at all? https://pytorch.org/blog/easily-list-and-initialize-models-with-new-apis-in-torchvision/

As much as I don't love the enum direction torchvision is going, I think we should continue to follow their API as closely as possible. We can always add a get_weight function that returns the right enum based on default choices. I say we go with 1, 2, 3 for now and add as many existing pretrained model weights as we can.