Overwriting FrozenBN eps=0.0 if pretrained=True for detection models.

test/test_models.py

@@ -8,7 +8,7 @@
 import unittest
 import random
 
-from torchvision.ops.misc import FrozenBatchNorm2d
+from torchvision.models._utils import overwrite_eps
 
 
 def set_rng_seed(seed):
@@ -151,9 +151,7 @@ def _test_detection_model(self, name, dev):
             kwargs["score_thresh"] = 0.013
         model = models.detection.__dict__[name](num_classes=50, pretrained_backbone=False, **kwargs)
         if "keypointrcnn" in name or "retinanet" in name:
-            for module in model.modules():
-                if isinstance(module, FrozenBatchNorm2d):
-                    module.eps = 0
+            overwrite_eps(model, 0.0)
         model.eval().to(device=dev)
         input_shape = (3, 300, 300)
         # RNG always on CPU, to ensure x in cuda tests is bitwise identical to x in cpu tests

torchvision/models/_utils.py

@@ -4,6 +4,8 @@
 from torch import nn
 from typing import Dict
 
+from ..ops.misc import FrozenBatchNorm2d
+
 
 class IntermediateLayerGetter(nn.ModuleDict):
     """
@@ -65,3 +67,9 @@ def forward(self, x):
                 out_name = self.return_layers[name]
                 out[out_name] = x
         return out
+
+
+def overwrite_eps(model, eps):
+    for module in model.modules():
+        if isinstance(module, FrozenBatchNorm2d):
+            module.eps = eps

torchvision/models/detection/faster_rcnn.py

@@ -7,6 +7,7 @@
 from torchvision.ops import misc as misc_nn_ops
 from torchvision.ops import MultiScaleRoIAlign
 
+from .._utils import overwrite_eps
 from ..utils import load_state_dict_from_url
 
 from .anchor_utils import AnchorGenerator
@@ -361,4 +362,5 @@ def fasterrcnn_resnet50_fpn(pretrained=False, progress=True,
         state_dict = load_state_dict_from_url(model_urls['fasterrcnn_resnet50_fpn_coco'],
                                               progress=progress)
         model.load_state_dict(state_dict)
+        overwrite_eps(model, 0.0)
     return model

torchvision/models/detection/keypoint_rcnn.py

@@ -3,6 +3,7 @@
 
 from torchvision.ops import MultiScaleRoIAlign
 
+from .._utils import overwrite_eps
 from ..utils import load_state_dict_from_url
 
 from .faster_rcnn import FasterRCNN
@@ -332,4 +333,5 @@ def keypointrcnn_resnet50_fpn(pretrained=False, progress=True,
         state_dict = load_state_dict_from_url(model_urls[key],
                                               progress=progress)
         model.load_state_dict(state_dict)
+        overwrite_eps(model, 0.0)
     return model

torchvision/models/detection/mask_rcnn.py

@@ -7,6 +7,7 @@
 from torchvision.ops import misc as misc_nn_ops
 from torchvision.ops import MultiScaleRoIAlign
 
+from .._utils import overwrite_eps
 from ..utils import load_state_dict_from_url
 
 from .faster_rcnn import FasterRCNN
@@ -328,4 +329,5 @@ def maskrcnn_resnet50_fpn(pretrained=False, progress=True,
         state_dict = load_state_dict_from_url(model_urls['maskrcnn_resnet50_fpn_coco'],
                                               progress=progress)
         model.load_state_dict(state_dict)
+        overwrite_eps(model, 0.0)
     return model

torchvision/models/detection/retinanet.py

@@ -7,6 +7,7 @@
 from torch import Tensor
 from torch.jit.annotations import Dict, List, Tuple
 
+from .._utils import overwrite_eps
 from ..utils import load_state_dict_from_url
 
 from . import _utils as det_utils
@@ -628,4 +629,5 @@ def retinanet_resnet50_fpn(pretrained=False, progress=True,
         state_dict = load_state_dict_from_url(model_urls['retinanet_resnet50_fpn_coco'],
                                               progress=progress)
         model.load_state_dict(state_dict)
+        overwrite_eps(model, 0.0)
     return model