Moving the method to detection utils and adding comments.

test/test_models.py

@@ -8,7 +8,7 @@
 import unittest
 import random
 
-from torchvision.models._utils import overwrite_eps
+from torchvision.models.detection._utils import overwrite_eps
 
 
 def set_rng_seed(seed):

torchvision/models/_utils.py

@@ -4,8 +4,6 @@
 from torch import nn
 from typing import Dict
 
-from ..ops.misc import FrozenBatchNorm2d
-
 
 class IntermediateLayerGetter(nn.ModuleDict):
     """
@@ -67,9 +65,3 @@ 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/_utils.py

@@ -3,7 +3,8 @@
 import torch
 from torch.jit.annotations import List, Tuple
 from torch import Tensor
-import torchvision
+
+from torchvision.ops.misc import FrozenBatchNorm2d
 
 
 class BalancedPositiveNegativeSampler(object):
@@ -349,3 +350,21 @@ def smooth_l1_loss(input, target, beta: float = 1. / 9, size_average: bool = Tru
     if size_average:
         return loss.mean()
     return loss.sum()
+
+
+def overwrite_eps(model, eps):
+    """
+    This method overwrites the default eps values of all the
+    FrozenBatchNorm2d layers of the model with the provided value.
+    This is necessary to address the BC-breaking change introduced
+    by the bug-fix at pytorch/vision#2933. The overwrite is applied
+    only when the pretrained weights are loaded to maintain compatibility
+    with previous versions.
+
+    Arguments:
+        model (nn.Module): The model on which we perform the overwrite.
+        eps (float): The new value of eps.
+    """
+    for module in model.modules():
+        if isinstance(module, FrozenBatchNorm2d):
+            module.eps = eps

torchvision/models/detection/faster_rcnn.py

@@ -7,7 +7,7 @@
 from torchvision.ops import misc as misc_nn_ops
 from torchvision.ops import MultiScaleRoIAlign
 
-from .._utils import overwrite_eps
+from ._utils import overwrite_eps
 from ..utils import load_state_dict_from_url
 
 from .anchor_utils import AnchorGenerator

torchvision/models/detection/keypoint_rcnn.py

@@ -3,7 +3,7 @@
 
 from torchvision.ops import MultiScaleRoIAlign
 
-from .._utils import overwrite_eps
+from ._utils import overwrite_eps
 from ..utils import load_state_dict_from_url
 
 from .faster_rcnn import FasterRCNN

torchvision/models/detection/mask_rcnn.py

@@ -7,7 +7,7 @@
 from torchvision.ops import misc as misc_nn_ops
 from torchvision.ops import MultiScaleRoIAlign
 
-from .._utils import overwrite_eps
+from ._utils import overwrite_eps
 from ..utils import load_state_dict_from_url
 
 from .faster_rcnn import FasterRCNN

torchvision/models/detection/retinanet.py

@@ -7,7 +7,7 @@
 from torch import Tensor
 from torch.jit.annotations import Dict, List, Tuple
 
-from .._utils import overwrite_eps
+from ._utils import overwrite_eps
 from ..utils import load_state_dict_from_url
 
 from . import _utils as det_utils