From 70131dedcbb39740729bf802986e814907cc62a9 Mon Sep 17 00:00:00 2001 From: AnikethCheluva Date: Fri, 23 Oct 2026 17:57:14 -0400 Subject: [PATCH] [Pipeline 9/12] make ResNet ImageNet normalization explicit and opt-in --- egomimic/models/stems/hpt_stems.py | 17 ++++++++++++++++- 1 file changed, 16 insertions(+), 1 deletion(-) diff --git a/egomimic/models/stems/hpt_stems.py b/egomimic/models/stems/hpt_stems.py index c99e71dd28..c4277b39c5 100644 --- a/egomimic/models/stems/hpt_stems.py +++ b/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 = []