[Pipeline 9/12] make ResNet ImageNet normalization explicit and opt-in

egomimic/models/stems/hpt_stems.py

@@ -163,9 +163,16 @@ def __init__(
         resnet_model: str = "resnet18",
         num_of_copy: int = 1,
         freeze_backbone: bool = False,
+        imagenet_normalize: bool = False,
         **kwargs,
     ) -> None:
-        """ResNet Encoder for Images"""
+        """ResNet Encoder for Images
+
+        ``imagenet_normalize`` applies the ImageNet channel statistics inside
+        :meth:`forward`, which is what pretrained ``weights`` expect. It defaults
+        to False because callers such as :class:`ResNetMLPImageStem` already normalize
+        outside the encoder; enabling it there would normalize twice.
+        """
         super().__init__(**kwargs)
         from egomimic.pipeline.construction import restoring_parameters
 
@@ -187,6 +194,11 @@ def __init__(
         self.input = input
         self.out_dim = output_dim
         self.to_tensor = transforms.ToTensor()
+        self.normalize = (
+            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
+            if imagenet_normalize
+            else None
+        )
         self.proj = nn.Linear(512, output_dim)
         self.avgpool = nn.AvgPool2d(7, stride=1)
 
@@ -215,6 +227,9 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         """
         B, *_, H, W = x.shape
         x = x.view(len(x), -1, 3, H, W)
+        if self.normalize is not None:
+            # images arrive in [0, 1]; pretrained weights expect ImageNet stats
+            x = self.normalize(x)
         if self.num_of_copy > 1:
             # separate encoding for each view
             out = []