add onnx support to roi_align and roi_pool (#593)

* add onnx support to roi_align and roi_pool

* add test on roi_align and roi_pool

* fix conflict with master

mmcv/ops/roi_align.py

+import torch
 import torch.nn as nn
 from torch.autograd import Function
 from torch.autograd.function import once_differentiable
@@ -14,16 +15,34 @@ class RoIAlignFunction(Function):
     @staticmethod
     def symbolic(g, input, rois, output_size, spatial_scale, sampling_ratio,
                  pool_mode, aligned):
+        from torch.onnx.symbolic_opset9 import sub, squeeze
+        from torch.onnx.symbolic_helper import _slice_helper
+        from torch.onnx import TensorProtoDataType
+        # batch_indices = rois[:, 0].long()
+        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 = rois[:, 1:]
+        rois = _slice_helper(g, rois, axes=[1], starts=[1], ends=[5])
+        if aligned:
+            # rois -= 0.5/spatial_scale
+            aligned_offset = g.op(
+                'Constant',
+                value_t=torch.tensor([0.5 / spatial_scale],
+                                     dtype=torch.float32))
+            rois = sub(g, rois, aligned_offset)
+        # roi align
         return g.op(
-            'MMCVRoIAlign',
+            'RoiAlign',
             input,
             rois,
-            aligned_height=output_size[0],
-            aligned_weight=output_size[1],
-            spatial_scale=spatial_scale,
-            sampling_ratio=sampling_ratio,
-            pool_mode=pool_mode,
-            aligned=aligned)
+            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,

mmcv/ops/roi_pool.py

@@ -15,12 +15,11 @@ class RoIPoolFunction(Function):
     @staticmethod
     def symbolic(g, input, rois, output_size, spatial_scale):
         return g.op(
-            'MMCVRoIPool',
+            'MaxRoiPool',
             input,
             rois,
-            pooled_height=output_size[0],
-            pooled_width=output_size[1],
-            spatial_scale=spatial_scale)
+            pooled_shape_i=output_size,
+            spatial_scale_f=spatial_scale)
 
     @staticmethod
     def forward(ctx, input, rois, output_size, spatial_scale=1.0):

tests/test_ops/test_onnx.py

@@ -56,3 +56,129 @@ def test_nms():
     onnx_score = onnx_dets[:, 4]
     os.remove(onnx_file)
     assert np.allclose(pytorch_score, onnx_score, atol=1e-3)
+
+
+def test_roialign():
+    from mmcv.ops import roi_align
+
+    # roi align config
+    pool_h = 2
+    pool_w = 2
+    spatial_scale = 1.0
+    sampling_ratio = 2
+
+    inputs = [([[[[1., 2.], [3., 4.]]]], [[0., 0., 0., 1., 1.]]),
+              ([[[[1., 2.], [3., 4.]], [[4., 3.],
+                                        [2., 1.]]]], [[0., 0., 0., 1., 1.]]),
+              ([[[[1., 2., 5., 6.], [3., 4., 7., 8.], [9., 10., 13., 14.],
+                  [11., 12., 15., 16.]]]], [[0., 0., 0., 3., 3.]])]
+
+    def warpped_function(torch_input, torch_rois):
+        return roi_align(torch_input, torch_rois, (pool_w, pool_h),
+                         spatial_scale, sampling_ratio, 'avg', True)
+
+    for case in inputs:
+        np_input = np.array(case[0], dtype=np.float32)
+        np_rois = np.array(case[1], dtype=np.float32)
+        input = torch.from_numpy(np_input)
+        rois = torch.from_numpy(np_rois)
+
+        # compute pytorch_output
+        with torch.no_grad():
+            pytorch_output = roi_align(input, rois, (pool_w, pool_h),
+                                       spatial_scale, sampling_ratio, 'avg',
+                                       True)
+
+        # export and load onnx model
+        wrapped_model = WrapFunction(warpped_function)
+        with torch.no_grad():
+            torch.onnx.export(
+                wrapped_model, (input, rois),
+                onnx_file,
+                export_params=True,
+                keep_initializers_as_inputs=True,
+                input_names=['input', 'rois'],
+                opset_version=11)
+        onnx_model = onnx.load(onnx_file)
+
+        # compute onnx_output
+        input_all = [node.name for node in onnx_model.graph.input]
+        input_initializer = [
+            node.name for node in onnx_model.graph.initializer
+        ]
+        net_feed_input = list(set(input_all) - set(input_initializer))
+        assert (len(net_feed_input) == 2)
+        sess = rt.InferenceSession(onnx_file)
+        onnx_output = sess.run(None, {
+            'input': input.detach().numpy(),
+            'rois': rois.detach().numpy()
+        })
+        onnx_output = onnx_output[0]
+
+        # allclose
+        os.remove(onnx_file)
+        assert np.allclose(pytorch_output, onnx_output, atol=1e-3)
+
+
+def test_roipool():
+    if not torch.cuda.is_available():
+        return
+    from mmcv.ops import roi_pool
+
+    # roi pool config
+    pool_h = 2
+    pool_w = 2
+    spatial_scale = 1.0
+
+    inputs = [([[[[1., 2.], [3., 4.]]]], [[0., 0., 0., 1., 1.]]),
+              ([[[[1., 2.], [3., 4.]], [[4., 3.],
+                                        [2., 1.]]]], [[0., 0., 0., 1., 1.]]),
+              ([[[[1., 2., 5., 6.], [3., 4., 7., 8.], [9., 10., 13., 14.],
+                  [11., 12., 15., 16.]]]], [[0., 0., 0., 3., 3.]])]
+
+    def warpped_function(torch_input, torch_rois):
+        return roi_pool(torch_input, torch_rois, (pool_w, pool_h),
+                        spatial_scale)
+
+    for case in inputs:
+        np_input = np.array(case[0], dtype=np.float32)
+        np_rois = np.array(case[1], dtype=np.float32)
+        input = torch.from_numpy(np_input).cuda()
+        rois = torch.from_numpy(np_rois).cuda()
+
+        # compute pytorch_output
+        with torch.no_grad():
+            pytorch_output = roi_pool(input, rois, (pool_w, pool_h),
+                                      spatial_scale)
+            pytorch_output = pytorch_output.cpu()
+
+        # export and load onnx model
+        wrapped_model = WrapFunction(warpped_function)
+        with torch.no_grad():
+            torch.onnx.export(
+                wrapped_model, (input, rois),
+                onnx_file,
+                export_params=True,
+                keep_initializers_as_inputs=True,
+                input_names=['input', 'rois'],
+                opset_version=11)
+        onnx_model = onnx.load(onnx_file)
+
+        # compute onnx_output
+        input_all = [node.name for node in onnx_model.graph.input]
+        input_initializer = [
+            node.name for node in onnx_model.graph.initializer
+        ]
+        net_feed_input = list(set(input_all) - set(input_initializer))
+        assert (len(net_feed_input) == 2)
+        sess = rt.InferenceSession(onnx_file)
+        onnx_output = sess.run(
+            None, {
+                'input': input.detach().cpu().numpy(),
+                'rois': rois.detach().cpu().numpy()
+            })
+        onnx_output = onnx_output[0]
+
+        # allclose
+        os.remove(onnx_file)
+        assert np.allclose(pytorch_output, onnx_output, atol=1e-3)