|
|
|
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 ext_loader
|
|
|
|
ext_module = ext_loader.load_ext('_ext',
|
|
['roi_pool_forward', 'roi_pool_backward'])
|
|
|
|
|
|
class RoIPoolFunction(Function):
|
|
|
|
@staticmethod
|
|
def symbolic(g, input, rois, output_size, spatial_scale):
|
|
return g.op(
|
|
'MaxRoiPool',
|
|
input,
|
|
rois,
|
|
pooled_shape_i=output_size,
|
|
spatial_scale_f=spatial_scale)
|
|
|
|
@staticmethod
|
|
def forward(ctx, input, rois, output_size, spatial_scale=1.0):
|
|
ctx.output_size = _pair(output_size)
|
|
ctx.spatial_scale = spatial_scale
|
|
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)
|
|
argmax = input.new_zeros(output_shape, dtype=torch.int)
|
|
|
|
ext_module.roi_pool_forward(
|
|
input,
|
|
rois,
|
|
output,
|
|
argmax,
|
|
pooled_height=ctx.output_size[0],
|
|
pooled_width=ctx.output_size[1],
|
|
spatial_scale=ctx.spatial_scale)
|
|
|
|
ctx.save_for_backward(rois, argmax)
|
|
return output
|
|
|
|
@staticmethod
|
|
@once_differentiable
|
|
def backward(ctx, grad_output):
|
|
rois, argmax = ctx.saved_tensors
|
|
grad_input = grad_output.new_zeros(ctx.input_shape)
|
|
|
|
ext_module.roi_pool_backward(
|
|
grad_output,
|
|
rois,
|
|
argmax,
|
|
grad_input,
|
|
pooled_height=ctx.output_size[0],
|
|
pooled_width=ctx.output_size[1],
|
|
spatial_scale=ctx.spatial_scale)
|
|
|
|
return grad_input, None, None, None
|
|
|
|
|
|
roi_pool = RoIPoolFunction.apply
|
|
|
|
|
|
class RoIPool(nn.Module):
|
|
|
|
def __init__(self, output_size, spatial_scale=1.0):
|
|
super(RoIPool, self).__init__()
|
|
|
|
self.output_size = _pair(output_size)
|
|
self.spatial_scale = float(spatial_scale)
|
|
|
|
def forward(self, input, rois):
|
|
return roi_pool(input, rois, self.output_size, self.spatial_scale)
|
|
|
|
def __repr__(self):
|
|
s = self.__class__.__name__
|
|
s += f'(output_size={self.output_size}, '
|
|
s += f'spatial_scale={self.spatial_scale})'
|
|
return s
|
|
|
