Repository: microsoft/torchgeo

Issue:
Improve model cards for pretrained weights

With the new v0.4 update to TorchGeo we have the ability to easily add pre-trained model weights (e.g. https://github.com/microsoft/torchgeo/blob/main/torchgeo/models/resnet.py#L28). Currently, each set of weights comes with an [`AugmentationSequential`](https://github.com/microsoft/torchgeo/blob/811a12c105ad839e05343ee243804b3694d36816/torchgeo/transforms/transforms.py#L19) instance that describes the preprocessing that the model expects, but don't include the proper normalization values (or a description of what range of values the model expects as input). These need to be retroactively added for all existing weights, and, ideally, verified (by trying to match reported downstream task performance with the model weights).

For example, the Sentinel 2 models from the SSL4EO papers likely expect uint12 Sentinel 2 values that are divided by 10,000 -- https://github.com/zhu-xlab/SSL4EO-S12/blob/main/src/benchmark/transfer_classification/datasets/EuroSat/eurosat_dataset.py#L97. This step should be added to the  [current set of transforms](https://github.com/microsoft/torchgeo/blob/main/torchgeo/models/resnet.py#L18) and tested.

If verification of the model weights is done on downstream tasks, we should also report those performance numbers vs. relying on the numbers reported in the associated papers.


Discussion:
I'd find this to be quite useful too! Found this issue since I was digging through code to figure out what normalization strategy to use for the SSL4EO pretrained models.
It seems like this is a popular idea, see https://github.com/zhu-xlab/SSL4EO-S12/issues/2 😉!

Resolving patch:
diff --git a/torchgeo/models/resnet.py b/torchgeo/models/resnet.py
--- a/torchgeo/models/resnet.py
+++ b/torchgeo/models/resnet.py
@@ -7,7 +7,7 @@
 
 import kornia.augmentation as K
 import timm
-import torch.nn as nn
+import torch
 from timm.models import ResNet
 from torchvision.models._api import Weights, WeightsEnum
 
@@ -15,8 +15,55 @@
 
 __all__ = ["ResNet50_Weights", "ResNet18_Weights"]
 
+
+# https://github.com/zhu-xlab/SSL4EO-S12/blob/d2868adfada65e40910bfcedfc49bc3b20df2248/src/benchmark/transfer_classification/linear_BE_moco.py#L167 # noqa: E501
+# https://github.com/zhu-xlab/SSL4EO-S12/blob/d2868adfada65e40910bfcedfc49bc3b20df2248/src/benchmark/transfer_classification/datasets/EuroSat/eurosat_dataset.py#L97 # noqa: E501
+# Normalization either by 10K or channel-wise with band statistics
 _zhu_xlab_transforms = AugmentationSequential(
-    K.Resize(256), K.CenterCrop(224), data_keys=["image"]
+    K.Resize(256),
+    K.CenterCrop(224),
+    K.Normalize(mean=0, std=10000),
+    data_keys=["image"],
+)
+
+# https://github.com/ServiceNow/seasonal-contrast/blob/8285173ec205b64bc3e53b880344dd6c3f79fa7a/datasets/bigearthnet_dataset.py#L13 # noqa: E501
+_seco_transforms = AugmentationSequential(
+    K.Resize(128),
+    K.Normalize(
+        mean=torch.Tensor(
+            [
+                340.76769064,
+                429.9430203,
+                614.21682446,
+                590.23569706,
+                950.68368468,
+                1792.46290469,
+                2075.46795189,
+                2218.94553375,
+                2266.46036911,
+                2246.0605464,
+                1594.42694882,
+                1009.32729131,
+            ]
+        ),
+        std=torch.Tensor(
+            [
+                554.81258967,
+                572.41639287,
+                582.87945694,
+                675.88746967,
+                729.89827633,
+                1096.01480586,
+                1273.45393088,
+                1365.45589904,
+                1356.13789355,
+                1302.3292881,
+                1079.19066363,
+                818.86747235,
+            ]
+        ),
+    ),
+    data_keys=["image"],
 )
 
 # https://github.com/pytorch/vision/pull/6883
@@ -62,7 +109,7 @@ class ResNet18_Weights(WeightsEnum):  # type: ignore[misc]
 
     SENTINEL2_RGB_SECO = Weights(
         url="https://huggingface.co/torchgeo/resnet18_sentinel2_rgb_seco/resolve/main/resnet18_sentinel2_rgb_seco-9976a9cb.pth",  # noqa: E501
-        transforms=nn.Identity(),
+        transforms=_seco_transforms,
         meta={
             "dataset": "SeCo Dataset",
             "in_chans": 3,
@@ -137,7 +184,7 @@ class ResNet50_Weights(WeightsEnum):  # type: ignore[misc]
 
     SENTINEL2_RGB_SECO = Weights(
         url="https://huggingface.co/torchgeo/resnet50_sentinel2_rgb_seco/resolve/main/resnet50_sentinel2_rgb_seco-584035db.pth",  # noqa: E501
-        transforms=nn.Identity(),
+        transforms=_seco_transforms,
         meta={
             "dataset": "SeCo Dataset",
             "in_chans": 3,
diff --git a/torchgeo/models/vit.py b/torchgeo/models/vit.py
--- a/torchgeo/models/vit.py
+++ b/torchgeo/models/vit.py
@@ -14,8 +14,14 @@
 
 __all__ = ["ViTSmall16_Weights"]
 
+# https://github.com/zhu-xlab/SSL4EO-S12/blob/d2868adfada65e40910bfcedfc49bc3b20df2248/src/benchmark/transfer_classification/linear_BE_moco.py#L167 # noqa: E501
+# https://github.com/zhu-xlab/SSL4EO-S12/blob/d2868adfada65e40910bfcedfc49bc3b20df2248/src/benchmark/transfer_classification/datasets/EuroSat/eurosat_dataset.py#L97 # noqa: E501
+# Normalization either by 10K or channel-wise with band statistics
 _zhu_xlab_transforms = AugmentationSequential(
-    K.Resize(256), K.CenterCrop(224), data_keys=["image"]
+    K.Resize(256),
+    K.CenterCrop(224),
+    K.Normalize(mean=0, std=10000),
+    data_keys=["image"],
 )
 
 # https://github.com/pytorch/vision/pull/6883


Test patch: