|
|
|
import torch
|
|
import torch.nn as nn
|
|
from torch.autograd import Function
|
|
from torch.autograd.function import once_differentiable
|
|
from torch.nn.modules.utils import _pair
|
|
|
|
from ..utils import deprecated_api_warning, ext_loader
|
|
|
|
ext_module = ext_loader.load_ext('_ext',
|
|
['roi_align_forward', 'roi_align_backward'])
|
|
|
|
|
|
class RoIAlignFunction(Function):
|
|
|
|
@staticmethod
|
|
def symbolic(g, input, rois, output_size, spatial_scale, sampling_ratio,
|
|
pool_mode, aligned):
|
|
from ..onnx import is_custom_op_loaded
|
|
has_custom_op = is_custom_op_loaded()
|
|
if has_custom_op:
|
|
return g.op(
|
|
'mmcv::MMCVRoiAlign',
|
|
input,
|
|
rois,
|
|
output_height_i=output_size[0],
|
|
output_width_i=output_size[1],
|
|
spatial_scale_f=spatial_scale,
|
|
sampling_ratio_i=sampling_ratio,
|
|
mode_s=pool_mode,
|
|
aligned_i=aligned)
|
|
else:
|
|
from torch.onnx.symbolic_opset9 import sub, squeeze
|
|
from torch.onnx.symbolic_helper import _slice_helper
|
|
from torch.onnx import TensorProtoDataType
|
|
|
|
batch_indices = _slice_helper(
|
|
g, rois, axes=[1], starts=[0], ends=[1])
|
|
batch_indices = squeeze(g, batch_indices, 1)
|
|
batch_indices = g.op(
|
|
'Cast', batch_indices, to_i=TensorProtoDataType.INT64)
|
|
|
|
rois = _slice_helper(g, rois, axes=[1], starts=[1], ends=[5])
|
|
if aligned:
|
|
|
|
aligned_offset = g.op(
|
|
'Constant',
|
|
value_t=torch.tensor([0.5 / spatial_scale],
|
|
dtype=torch.float32))
|
|
rois = sub(g, rois, aligned_offset)
|
|
|
|
return g.op(
|
|
'RoiAlign',
|
|
input,
|
|
rois,
|
|
batch_indices,
|
|
output_height_i=output_size[0],
|
|
output_width_i=output_size[1],
|
|
spatial_scale_f=spatial_scale,
|
|
sampling_ratio_i=max(0, sampling_ratio),
|
|
mode_s=pool_mode)
|
|
|
|
@staticmethod
|
|
def forward(ctx,
|
|
input,
|
|
rois,
|
|
output_size,
|
|
spatial_scale=1.0,
|
|
sampling_ratio=0,
|
|
pool_mode='avg',
|
|
aligned=True):
|
|
ctx.output_size = _pair(output_size)
|
|
ctx.spatial_scale = spatial_scale
|
|
ctx.sampling_ratio = sampling_ratio
|
|
assert pool_mode in ('max', 'avg')
|
|
ctx.pool_mode = 0 if pool_mode == 'max' else 1
|
|
ctx.aligned = aligned
|
|
ctx.input_shape = input.size()
|
|
|
|
assert rois.size(1) == 5, 'RoI must be (idx, x1, y1, x2, y2)!'
|
|
|
|
output_shape = (rois.size(0), input.size(1), ctx.output_size[0],
|
|
ctx.output_size[1])
|
|
output = input.new_zeros(output_shape)
|
|
if ctx.pool_mode == 0:
|
|
argmax_y = input.new_zeros(output_shape)
|
|
argmax_x = input.new_zeros(output_shape)
|
|
else:
|
|
argmax_y = input.new_zeros(0)
|
|
argmax_x = input.new_zeros(0)
|
|
|
|
ext_module.roi_align_forward(
|
|
input,
|
|
rois,
|
|
output,
|
|
argmax_y,
|
|
argmax_x,
|
|
aligned_height=ctx.output_size[0],
|
|
aligned_width=ctx.output_size[1],
|
|
spatial_scale=ctx.spatial_scale,
|
|
sampling_ratio=ctx.sampling_ratio,
|
|
pool_mode=ctx.pool_mode,
|
|
aligned=ctx.aligned)
|
|
|
|
ctx.save_for_backward(rois, argmax_y, argmax_x)
|
|
return output
|
|
|
|
@staticmethod
|
|
@once_differentiable
|
|
def backward(ctx, grad_output):
|
|
rois, argmax_y, argmax_x = ctx.saved_tensors
|
|
grad_input = grad_output.new_zeros(ctx.input_shape)
|
|
|
|
grad_output = grad_output.contiguous()
|
|
ext_module.roi_align_backward(
|
|
grad_output,
|
|
rois,
|
|
argmax_y,
|
|
argmax_x,
|
|
grad_input,
|
|
aligned_height=ctx.output_size[0],
|
|
aligned_width=ctx.output_size[1],
|
|
spatial_scale=ctx.spatial_scale,
|
|
sampling_ratio=ctx.sampling_ratio,
|
|
pool_mode=ctx.pool_mode,
|
|
aligned=ctx.aligned)
|
|
return grad_input, None, None, None, None, None, None
|
|
|
|
|
|
roi_align = RoIAlignFunction.apply
|
|
|
|
|
|
class RoIAlign(nn.Module):
|
|
"""RoI align pooling layer.
|
|
|
|
Args:
|
|
output_size (tuple): h, w
|
|
spatial_scale (float): scale the input boxes by this number
|
|
sampling_ratio (int): number of inputs samples to take for each
|
|
output sample. 0 to take samples densely for current models.
|
|
pool_mode (str, 'avg' or 'max'): pooling mode in each bin.
|
|
aligned (bool): if False, use the legacy implementation in
|
|
MMDetection. If True, align the results more perfectly.
|
|
use_torchvision (bool): whether to use roi_align from torchvision.
|
|
|
|
Note:
|
|
The implementation of RoIAlign when aligned=True is modified from
|
|
https://github.com/facebookresearch/detectron2/
|
|
|
|
The meaning of aligned=True:
|
|
|
|
Given a continuous coordinate c, its two neighboring pixel
|
|
indices (in our pixel model) are computed by floor(c - 0.5) and
|
|
ceil(c - 0.5). For example, c=1.3 has pixel neighbors with discrete
|
|
indices [0] and [1] (which are sampled from the underlying signal
|
|
at continuous coordinates 0.5 and 1.5). But the original roi_align
|
|
(aligned=False) does not subtract the 0.5 when computing
|
|
neighboring pixel indices and therefore it uses pixels with a
|
|
slightly incorrect alignment (relative to our pixel model) when
|
|
performing bilinear interpolation.
|
|
|
|
With `aligned=True`,
|
|
we first appropriately scale the ROI and then shift it by -0.5
|
|
prior to calling roi_align. This produces the correct neighbors;
|
|
|
|
The difference does not make a difference to the model's
|
|
performance if ROIAlign is used together with conv layers.
|
|
"""
|
|
|
|
@deprecated_api_warning(
|
|
{
|
|
'out_size': 'output_size',
|
|
'sample_num': 'sampling_ratio'
|
|
},
|
|
cls_name='RoIAlign')
|
|
def __init__(self,
|
|
output_size,
|
|
spatial_scale=1.0,
|
|
sampling_ratio=0,
|
|
pool_mode='avg',
|
|
aligned=True,
|
|
use_torchvision=False):
|
|
super(RoIAlign, self).__init__()
|
|
|
|
self.output_size = _pair(output_size)
|
|
self.spatial_scale = float(spatial_scale)
|
|
self.sampling_ratio = int(sampling_ratio)
|
|
self.pool_mode = pool_mode
|
|
self.aligned = aligned
|
|
self.use_torchvision = use_torchvision
|
|
|
|
def forward(self, input, rois):
|
|
"""
|
|
Args:
|
|
input: NCHW images
|
|
rois: Bx5 boxes. First column is the index into N.\
|
|
The other 4 columns are xyxy.
|
|
"""
|
|
if self.use_torchvision:
|
|
from torchvision.ops import roi_align as tv_roi_align
|
|
if 'aligned' in tv_roi_align.__code__.co_varnames:
|
|
return tv_roi_align(input, rois, self.output_size,
|
|
self.spatial_scale, self.sampling_ratio,
|
|
self.aligned)
|
|
else:
|
|
if self.aligned:
|
|
rois -= rois.new_tensor([0.] +
|
|
[0.5 / self.spatial_scale] * 4)
|
|
return tv_roi_align(input, rois, self.output_size,
|
|
self.spatial_scale, self.sampling_ratio)
|
|
else:
|
|
return roi_align(input, rois, self.output_size, self.spatial_scale,
|
|
self.sampling_ratio, self.pool_mode, self.aligned)
|
|
|
|
def __repr__(self):
|
|
s = self.__class__.__name__
|
|
s += f'(output_size={self.output_size}, '
|
|
s += f'spatial_scale={self.spatial_scale}, '
|
|
s += f'sampling_ratio={self.sampling_ratio}, '
|
|
s += f'pool_mode={self.pool_mode}, '
|
|
s += f'aligned={self.aligned}, '
|
|
s += f'use_torchvision={self.use_torchvision})'
|
|
return s
|
|
|
