[v1.x] Onnx export support for ROIAlign

python/mxnet/contrib/onnx/mx2onnx/_op_translations.py

@@ -3555,3 +3555,42 @@ def convert_broadcast_like(node, **kwargs):
     ]
 
     return nodes
+
+
+@mx_op.register('_contrib_ROIAlign')
+def convert_contrib_roialign(node, **kwargs):
+    """Map MXNet's _contrib_ROIAlign
+    """
+    from onnx.helper import make_node
+    from onnx import TensorProto
+    name, input_nodes, attrs = get_inputs(node, kwargs)
+
+    pooled_size = convert_string_to_list(str(attrs.get('pooled_size')))
+    spatial_scale = float(attrs.get('spatial_scale'))
+    sample_ratio = int(attrs.get('sample_ratio', '0'))
+    position_sensitive = attrs.get('position_sensitive', 'False')
+    aligned = attrs.get('aligned', 'False')
+
+    '''
+    if position_sensitive != 'False':
+        raise NotImplementedError('_contrib_ROIAlign does not currently support \
+                                   position_sensitive!=False')
+    if aligned != 'False':
+        raise NotImplementedError('_contrib_ROIAlign does not currently support \
+                                   aligned!=False')
+    '''
+    create_tensor([0], name+'_0', kwargs['initializer']),
+    create_tensor([1], name+'_1', kwargs['initializer']),
+    create_tensor([5], name+'_5', kwargs['initializer']),
+
+    nodes = [
+        make_node('Slice', [input_nodes[1], name+'_1', name+'_5', name+'_1'], [name+'_rois']),
+        make_node('Slice', [input_nodes[1], name+'_0', name+'_1', name+'_1'], [name+'_inds__']),
+        make_node('Squeeze', [name+'_inds__'], [name+'_inds_'], axes=(1,)),
+        make_node('Cast', [name+'_inds_'], [name+'_inds'], to=int(TensorProto.INT64)),
+        make_node('RoiAlign', [input_nodes[0], name+'_rois', name+'_inds'], [name],
+                  mode='avg', output_height=pooled_size[0], output_width=pooled_size[1],
+                  sampling_ratio=sample_ratio, spatial_scale=spatial_scale)
+    ]
+
+    return nodes

tests/python-pytest/onnx/test_operators.py

@@ -644,3 +644,20 @@ def test_onnx_export_broadcast_like(tmp_path, dtype, lhs_axes, rhs_axes):
     op_export_test('broadcast_like1', M1, [x, y], tmp_path)
     M2 = def_model('broadcast_like', lhs_axes=lhs_axes, rhs_axes=rhs_axes)
     op_export_test('broadcast_like2', M2, [x, y], tmp_path)
+
+
+@pytest.mark.parametrize('dtype', ['float32'])
+@pytest.mark.parametrize('pooled_size', [(1, 1), (3, 3), (14, 14), (5, 7)])
+@pytest.mark.parametrize('spatial_scale', [1, 0.5, 0.0625])
+@pytest.mark.parametrize('spatial_ratio', [1, 2, 3, 5])
+def test_onnx_export_contrib_ROIAlign(tmp_path, dtype, pooled_size, spatial_scale, spatial_ratio):
+    data = mx.random.uniform(0, 1, (5, 3, 128, 128)).astype(dtype)
+    rois = mx.nd.array([[0, 0, 0, 63, 63],
+                        [1, 34, 52, 25, 85],
+                        [2, 50, 50, 100, 100],
+                        [3, 0, 0, 127, 127],
+                        [4, 12, 84, 22, 94],
+                        [0, 0, 0, 1, 1]]).astype(dtype)
+    M = def_model('contrib.ROIAlign', pooled_size=pooled_size, spatial_scale=spatial_scale,
+                  sample_ratio=spatial_ratio)
+    op_export_test('_contrib_ROIAlign', M, [data, rois], tmp_path)