use torchvision's roialign
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
from torch import nn
+from torchvision.ops import roi_align as tv_roi_align
+
+try:
+ from torchvision import __version__
+
+ version = tuple(int(x) for x in __version__.split(".")[:2])
+ USE_TORCHVISION = version >= (0, 7) # https://github.com/pytorch/vision/pull/2438
+except ImportError: # only open source torchvision has __version__
+ USE_TORCHVISION = True
+
+
+if USE_TORCHVISION:
+ roi_align = tv_roi_align
+else:
+ from torch.nn.modules.utils import _pair
from torch.autograd import Function
from torch.autograd.function import once_differentiable
-from torch.nn.modules.utils import _pair
-
from detectron2 import _C
-
class _ROIAlign(Function):
@staticmethod
def forward(ctx, input, roi, output_size, spatial_scale, sampling_ratio, aligned):
@@ -44,10 +56,10 @@ class _ROIAlign(Function):
)
return grad_input, None, None, None, None, None
-
roi_align = _ROIAlign.apply
+# NOTE: torchvision's RoIAlign has a different default aligned=False
class ROIAlign(nn.Module):
def __init__(self, output_size, spatial_scale, sampling_ratio, aligned=True):
"""
