| |
| 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 |
|
|
