Use crop_and_resize from longcw/RoIAlign.pytorch. Speed up the training speed from 0.315s -> 0.282 (on my machine).

Author: ruotianluo

README.md

@@ -69,7 +69,7 @@ Additional features not mentioned in the [report](https://arxiv.org/pdf/1702.021
   git clone https://github.com/ruotianluo/pytorch-faster-rcnn.git
   ```
 
-2. Compile modules:
+2. Compile modules(nms, roi_pooling(from [longcw/yolo2-pytorch](https://github.com/longcw/yolo2-pytorch)), roi_align(from [longcw/RoIAlign.pytorch](https://github.com/longcw/RoIAlign.pytorch.git))):
   ```
   cd pytorch-faster-rcnn/lib
   bash make.sh

lib/layer_utils/roi_align/__init__.py

@@ -0,0 +1,15 @@
+
+from torch.utils.ffi import _wrap_function
+from ._crop_and_resize import lib as _lib, ffi as _ffi
+
+__all__ = []
+def _import_symbols(locals):
+    for symbol in dir(_lib):
+        fn = getattr(_lib, symbol)
+        if callable(fn):
+            locals[symbol] = _wrap_function(fn, _ffi)
+        else:
+            locals[symbol] = fn
+        __all__.append(symbol)
+
+_import_symbols(locals())

lib/layer_utils/roi_align/_ext/__init__.py

@@ -0,0 +1,40 @@
+import os
+import torch
+from torch.utils.ffi import create_extension
+
+
+sources = ['src/crop_and_resize.c']
+headers = ['src/crop_and_resize.h']
+defines = []
+with_cuda = False
+
+extra_objects = []
+if torch.cuda.is_available():
+    print('Including CUDA code.')
+    sources += ['src/crop_and_resize_gpu.c']
+    headers += ['src/crop_and_resize_gpu.h']
+    defines += [('WITH_CUDA', None)]
+    extra_objects += ['src/cuda/crop_and_resize_kernel.cu.o']
+    with_cuda = True
+
+extra_compile_args = ['-std=c99']
+
+this_file = os.path.dirname(os.path.realpath(__file__))
+print(this_file)
+sources = [os.path.join(this_file, fname) for fname in sources]
+headers = [os.path.join(this_file, fname) for fname in headers]
+extra_objects = [os.path.join(this_file, fname) for fname in extra_objects]
+
+ffi = create_extension(
+    '_ext.crop_and_resize',
+    headers=headers,
+    sources=sources,
+    define_macros=defines,
+    relative_to=__file__,
+    with_cuda=with_cuda,
+    extra_objects=extra_objects,
+    extra_compile_args=extra_compile_args
+)
+
+if __name__ == '__main__':
+    ffi.build()

lib/layer_utils/roi_align/_ext/crop_and_resize/__init__.py

@@ -0,0 +1,67 @@
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.autograd import Function
+
+from ._ext import crop_and_resize as _backend
+
+
+class CropAndResizeFunction(Function):
+
+    def __init__(self, crop_height, crop_width, extrapolation_value=0):
+        self.crop_height = crop_height
+        self.crop_width = crop_width
+        self.extrapolation_value = extrapolation_value
+
+    def forward(self, image, boxes, box_ind):
+        crops = torch.zeros_like(image)
+
+        if image.is_cuda:
+            _backend.crop_and_resize_gpu_forward(
+                image, boxes, box_ind,
+                self.extrapolation_value, self.crop_height, self.crop_width, crops)
+        else:
+            _backend.crop_and_resize_forward(
+                image, boxes, box_ind,
+                self.extrapolation_value, self.crop_height, self.crop_width, crops)
+
+        # save for backward
+        self.im_size = image.size()
+        self.save_for_backward(boxes, box_ind)
+
+        return crops
+
+    def backward(self, grad_outputs):
+        boxes, box_ind = self.saved_tensors
+
+        grad_outputs = grad_outputs.contiguous()
+        grad_image = torch.zeros_like(grad_outputs).resize_(*self.im_size)
+
+        if grad_outputs.is_cuda:
+            _backend.crop_and_resize_gpu_backward(
+                grad_outputs, boxes, box_ind, grad_image
+            )
+        else:
+            _backend.crop_and_resize_backward(
+                grad_outputs, boxes, box_ind, grad_image
+            )
+
+        return grad_image, None, None
+
+
+class CropAndResize(nn.Module):
+    """
+    Crop and resize ported from tensorflow
+    See more details on https://www.tensorflow.org/api_docs/python/tf/image/crop_and_resize
+    """
+
+    def __init__(self, crop_height, crop_width, extrapolation_value=0):
+        super(CropAndResize, self).__init__()
+
+        self.crop_height = crop_height
+        self.crop_width = crop_width
+        self.extrapolation_value = extrapolation_value
+
+    def forward(self, image, boxes, box_ind):
+        return CropAndResizeFunction(self.crop_height, self.crop_width, self.extrapolation_value)(image, boxes, box_ind)

lib/layer_utils/roi_align/build.py

@@ -0,0 +1,48 @@
+import torch
+from torch import nn
+
+from .crop_and_resize import CropAndResizeFunction, CropAndResize
+
+
+class RoIAlign(nn.Module):
+
+    def __init__(self, crop_height, crop_width, extrapolation_value=0, transform_fpcoor=True):
+        super(RoIAlign, self).__init__()
+
+        self.crop_height = crop_height
+        self.crop_width = crop_width
+        self.extrapolation_value = extrapolation_value
+        self.transform_fpcoor = transform_fpcoor
+
+    def forward(self, featuremap, boxes, box_ind):
+        """
+        RoIAlign based on crop_and_resize.
+        See more details on https://github.com/ppwwyyxx/tensorpack/blob/6d5ba6a970710eaaa14b89d24aace179eb8ee1af/examples/FasterRCNN/model.py#L301
+        :param featuremap: NxCxHxW
+        :param boxes: Mx4 float box with (x1, y1, x2, y2) **without normalization**
+        :param box_ind: M
+        :return: MxCxoHxoW
+        """
+        x1, y1, x2, y2 = torch.split(boxes, 1, dim=1)
+        image_height, image_width = featuremap.size()[2:4]
+
+        if self.transform_fpcoor:
+            spacing_w = (x2 - x1) / float(self.crop_width)
+            spacing_h = (y2 - y1) / float(self.crop_height)
+
+            nx0 = (x1 + spacing_w / 2 - 0.5) / float(image_width - 1)
+            ny0 = (y1 + spacing_h / 2 - 0.5) / float(image_height - 1)
+            nw = spacing_w * float(self.crop_width - 1) / float(image_width - 1)
+            nh = spacing_h * float(self.crop_height - 1) / float(image_height - 1)
+
+            boxes = torch.cat((ny0, nx0, ny0 + nh, nx0 + nw), 1)
+        else:
+            x1 = x1 / float(image_width - 1)
+            x2 = x2 / float(image_width - 1)
+            y1 = y1 / float(image_height - 1)
+            y2 = y2 / float(image_height - 1)
+            boxes = torch.cat((y1, x1, y2, x2), 1)
+
+        boxes = boxes.detach().contiguous()
+        box_ind = box_ind.detach()
+        return CropAndResizeFunction(self.crop_height, self.crop_width, self.extrapolation_value)(featuremap, boxes, box_ind)