Scale-MAE model

[isaaccorley]

Finally getting around to adding this.

Adds Scale-MAE model (ViT encoder only) and pretrained weights.

I've verified this reproduces KNN performance at different resolutions for UCMerced but will repeat for other datasets.

@RitwikGupta let me know if this looks good. I cleaned up some of the code a bit to work out of the box with our trainers (this required setting the res when initializing the model instead of dynamically but I think it should still be fine).

@calebrob6 lmk if you want to team up on this one.

[adamjstewart]

Also see changes in https://github.com/microsoft/torchgeo/pull/2052/files

[adamjstewart]

Some minor renaming suggestions to make things more consistent with DOFA, and some major documentation improvement suggestions. I'm willing to help with both if needed.

[adamjstewart]

Two thoughts:

1. I wonder if we should move both of these to the new "Sensor-Agnostic" section because they technically work for (RGB-only) imagery from any sensor
2. Since both of these have evaluation results on fMoW, can we add additional columns with those performance metrics (assuming they are comparable)? If we move them to "Sensor-Agnostic", we may need two tables, one for things evaluated on GEO-Bench and one for things evaluated on fMoW.

[adamjstewart]

We may also want to add a short summary or table of which sensor-agnostic models provide which features. For example, DOFA enables explicit dynamic spectral band support (via model arch) and implicit dynamic resolution (via training data) while Scale-MAE has no dynamic spectral band support (RGB-only) but explicit dynamic resolution support (via model arch). Not sure about GASSL, maybe only implicit dynamic resolution (via training data)? It's worth mentioning that neither have dynamic temporal resolution support (maybe Satlas does?). I'm planning on highlighting this in our release notes, so I can also write something up if needed. Something like:

"The following pre-trained models offer dynamic spatial (resolution), temporal (time span), and/or spectral (wavelength) support, either via their training data (implicit) or via their model architecture (explicit):"

Model	Spatial	Temporal	Spectral
DOFA	implicit	-	explicit
GASSL	implicit	-	-
Scale-MAE	explicit	-	-

We could also optionally specify the range of resolutions/time spans/wavelengths that the model was pre-trained on. Just want to give users more feedback as to which model to choose.

[isaaccorley]

Should we save this for a different PR?

[adamjstewart]

I'm fine with that, just don't let me forget before the release.

[isaaccorley]

I'll open a PR after this one so we don't forget to finish it

[adamjstewart]

Suggested change

	Scale-MAE Vision Transformer
	^^^^^^^^^^^^^^^^^^^^^^^^^^^^
	Scale-MAE
	^^^^^^^^^

This is how we named DOFA, which also uses a ViT backbone

[adamjstewart]

Suggested change

	"""Pre-trained Scale-MAE Vision Transformer models."""
	"""Pre-trained Scale-MAE models."""

[adamjstewart]

Suggested change

	class ScaleMAEViT(VisionTransformer): # type: ignore[misc]
	class ScaleMAE(VisionTransformer): # type: ignore[misc]

[adamjstewart]

Suggested change

	class ScaleMAE_ViTLarge16_Weights(WeightsEnum): # type: ignore[misc]
	class ScaleMAELarge16_Weights(WeightsEnum): # type: ignore[misc]

[adamjstewart]

Suggested change

	def scalemae_vit_large_patch16(
	def scalemae_large_patch16(

[adamjstewart]

Or is the image size customizable?

[adamjstewart]

Can we add additional sizes? I know we only have pre-trained weights available for large, but we might as well add functions to instantiate other sizes like they do in the source repo.

[isaaccorley]

The image size can be changed and the positional embeddings will get interpolated.

[adamjstewart]

Updated diff to remove 244