# Copyright (c) OpenMMLab. All rights reserved. import os import warnings from functools import partial import numpy as np import onnx import onnxruntime as rt import pytest import torch import torch.nn as nn import torch.nn.functional as F from packaging import version onnx_file = 'tmp.onnx' if torch.__version__ == 'parrots': pytest.skip('not supported in parrots now', allow_module_level=True) @pytest.fixture(autouse=True) def run_before_and_after_test(): # clear onnx_file before test if os.path.exists(onnx_file): os.remove(onnx_file) yield # clear onnx_file after test if os.path.exists(onnx_file): os.remove(onnx_file) class WrapFunction(nn.Module): def __init__(self, wrapped_function): super().__init__() self.wrapped_function = wrapped_function def forward(self, *args, **kwargs): return self.wrapped_function(*args, **kwargs) def process_grid_sample(func, input, grid, ort_custom_op_path=''): wrapped_model = WrapFunction(func).eval() input_names = ['input', 'grid'] output_names = ['output'] with torch.no_grad(): torch.onnx.export( wrapped_model, (input, grid), onnx_file, export_params=True, keep_initializers_as_inputs=True, input_names=input_names, output_names=output_names, opset_version=11) onnx_model = onnx.load(onnx_file) session_options = rt.SessionOptions() if ort_custom_op_path: session_options.register_custom_ops_library(ort_custom_op_path) # get 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, session_options) ort_result = sess.run(None, { 'input': input.detach().numpy(), 'grid': grid.detach().numpy() }) pytorch_results = wrapped_model(input.clone(), grid.clone()) assert np.allclose(pytorch_results, ort_result, atol=1e-3) @pytest.mark.parametrize('mode', ['bilinear', 'nearest']) @pytest.mark.parametrize('padding_mode', ['zeros', 'border', 'reflection']) @pytest.mark.parametrize('align_corners', [True, False]) def test_grid_sample(mode, padding_mode, align_corners): from mmcv.onnx.symbolic import register_extra_symbolics opset_version = 11 register_extra_symbolics(opset_version) from mmcv.ops import get_onnxruntime_op_path ort_custom_op_path = get_onnxruntime_op_path() if not os.path.exists(ort_custom_op_path): pytest.skip('custom ops for onnxruntime are not compiled.') input = torch.rand(1, 1, 10, 10) grid = torch.Tensor([[[1, 0, 0], [0, 1, 0]]]) grid = F.affine_grid( grid, (1, 1, 15, 15), align_corners=align_corners).type_as(input) def func(input, grid): return F.grid_sample( input, grid, mode=mode, padding_mode=padding_mode, align_corners=align_corners) return process_grid_sample(func, input, grid, ort_custom_op_path) @pytest.mark.parametrize('align_corners', [True, False]) def test_bilinear_grid_sample(align_corners): from mmcv.ops.point_sample import bilinear_grid_sample # only support pytorch >= 1.5.0 if version.parse(torch.__version__) < version.parse('1.5.0'): pytest.skip('Only support PyTorch >= 1.5.0') input = torch.rand(1, 1, 10, 10) grid = torch.Tensor([[[1, 0, 0], [0, 1, 0]]]) grid = F.affine_grid( grid, (1, 1, 15, 15), align_corners=align_corners).type_as(input) def func(input, grid): return bilinear_grid_sample(input, grid, align_corners=align_corners) return process_grid_sample(func, input, grid) def test_nms(): from mmcv.ops import get_onnxruntime_op_path, nms np_boxes = np.array([[6.0, 3.0, 8.0, 7.0], [3.0, 6.0, 9.0, 11.0], [3.0, 7.0, 10.0, 12.0], [1.0, 4.0, 13.0, 7.0]], dtype=np.float32) np_scores = np.array([0.6, 0.9, 0.7, 0.2], dtype=np.float32) boxes = torch.from_numpy(np_boxes) scores = torch.from_numpy(np_scores) nms = partial( nms, iou_threshold=0.3, offset=0, score_threshold=0, max_num=0) pytorch_dets, _ = nms(boxes, scores) pytorch_score = pytorch_dets[:, 4] wrapped_model = WrapFunction(nms) wrapped_model.cpu().eval() with torch.no_grad(): torch.onnx.export( wrapped_model, (boxes, scores), onnx_file, export_params=True, keep_initializers_as_inputs=True, input_names=['boxes', 'scores'], opset_version=11) onnx_model = onnx.load(onnx_file) ort_custom_op_path = get_onnxruntime_op_path() session_options = rt.SessionOptions() if os.path.exists(ort_custom_op_path): session_options.register_custom_ops_library(ort_custom_op_path) # get 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, session_options) onnx_dets, _ = sess.run(None, { 'scores': scores.detach().numpy(), 'boxes': boxes.detach().numpy() }) onnx_score = onnx_dets[:, 4] assert np.allclose(pytorch_score, onnx_score, atol=1e-3) @pytest.mark.skipif(not torch.cuda.is_available(), reason='test requires GPU') def test_softnms(): from mmcv.ops import get_onnxruntime_op_path, soft_nms # only support pytorch >= 1.7.0 if version.parse(torch.__version__) < version.parse('1.7.0'): warnings.warn('test_softnms should be ran with pytorch >= 1.7.0') return # only support onnxruntime >= 1.5.1 assert version.parse(rt.__version__) >= version.parse( '1.5.1'), 'test_softnms should be ran with onnxruntime >= 1.5.1' ort_custom_op_path = get_onnxruntime_op_path() if not os.path.exists(ort_custom_op_path): pytest.skip('softnms for onnxruntime is not compiled.') np_boxes = np.array([[6.0, 3.0, 8.0, 7.0], [3.0, 6.0, 9.0, 11.0], [3.0, 7.0, 10.0, 12.0], [1.0, 4.0, 13.0, 7.0]], dtype=np.float32) np_scores = np.array([0.6, 0.9, 0.7, 0.2], dtype=np.float32) boxes = torch.from_numpy(np_boxes) scores = torch.from_numpy(np_scores) configs = [[0.3, 0.5, 0.01, 'linear'], [0.3, 0.5, 0.01, 'gaussian'], [0.3, 0.5, 0.01, 'naive']] session_options = rt.SessionOptions() session_options.register_custom_ops_library(ort_custom_op_path) for _iou_threshold, _sigma, _min_score, _method in configs: pytorch_dets, pytorch_inds = soft_nms( boxes, scores, iou_threshold=_iou_threshold, sigma=_sigma, min_score=_min_score, method=_method) nms = partial( soft_nms, iou_threshold=_iou_threshold, sigma=_sigma, min_score=_min_score, method=_method) wrapped_model = WrapFunction(nms) wrapped_model.cpu().eval() with torch.no_grad(): torch.onnx.export( wrapped_model, (boxes, scores), onnx_file, export_params=True, keep_initializers_as_inputs=True, input_names=['boxes', 'scores'], opset_version=11) onnx_model = onnx.load(onnx_file) # get 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, session_options) onnx_dets, onnx_inds = sess.run(None, { 'scores': scores.detach().numpy(), 'boxes': boxes.detach().numpy() }) assert np.allclose(pytorch_dets, onnx_dets, atol=1e-3) assert np.allclose(onnx_inds, onnx_inds, atol=1e-3) def test_roialign(): try: from mmcv.ops import get_onnxruntime_op_path, roi_align except (ImportError, ModuleNotFoundError): pytest.skip('roi_align op is not successfully compiled') ort_custom_op_path = get_onnxruntime_op_path() # 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) session_options = rt.SessionOptions() if os.path.exists(ort_custom_op_path): session_options.register_custom_ops_library(ort_custom_op_path) # 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, session_options) onnx_output = sess.run(None, { 'input': input.detach().numpy(), 'rois': rois.detach().numpy() }) onnx_output = onnx_output[0] # allclose assert np.allclose(pytorch_output, onnx_output, atol=1e-3) def test_roialign_rotated(): try: from mmcv.ops import get_onnxruntime_op_path, roi_align_rotated except (ImportError, ModuleNotFoundError): pytest.skip('roi_align_aligned op is not successfully compiled') ort_custom_op_path = get_onnxruntime_op_path() if not os.path.exists(ort_custom_op_path): pytest.skip('custom ops for onnxruntime are not compiled.') # roi align config pool_h = 2 pool_w = 2 spatial_scale = 1.0 sampling_ratio = 2 inputs = [([[[[1., 2.], [3., 4.]]]], [[0., 0.5, 0.5, 1., 1., 0]]), ([[[[1., 2.], [3., 4.]]]], [[0., 0.5, 0.5, 1., 1., np.pi / 2]]), ([[[[1., 2.], [3., 4.]], [[4., 3.], [2., 1.]]]], [[0., 0.5, 0.5, 1., 1., 0]]), ([[[[1., 2., 5., 6.], [3., 4., 7., 8.], [9., 10., 13., 14.], [11., 12., 15., 16.]]]], [[0., 1.5, 1.5, 3., 3., 0]]), ([[[[1., 2., 5., 6.], [3., 4., 7., 8.], [9., 10., 13., 14.], [11., 12., 15., 16.]]]], [[0., 1.5, 1.5, 3., 3., np.pi / 2]])] def warpped_function(torch_input, torch_rois): return roi_align_rotated(torch_input, torch_rois, (pool_w, pool_h), spatial_scale, sampling_ratio, True, False) 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_rotated(input, rois, (pool_w, pool_h), spatial_scale, sampling_ratio, True, False) # 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=['features', 'rois'], opset_version=11) onnx_model = onnx.load(onnx_file) session_options = rt.SessionOptions() if os.path.exists(ort_custom_op_path): session_options.register_custom_ops_library(ort_custom_op_path) # 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, session_options) onnx_output = sess.run(None, { 'features': input.detach().numpy(), 'rois': rois.detach().numpy() }) onnx_output = onnx_output[0] # allclose assert np.allclose(pytorch_output, onnx_output, atol=1e-3) @pytest.mark.skipif(not torch.cuda.is_available(), reason='test requires GPU') def test_roipool(): 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 assert np.allclose(pytorch_output, onnx_output, atol=1e-3) def test_interpolate(): from mmcv.onnx.symbolic import register_extra_symbolics opset_version = 11 register_extra_symbolics(opset_version) def func(feat, scale_factor=2): out = F.interpolate(feat, scale_factor=scale_factor) return out net = WrapFunction(func) net = net.cpu().eval() dummy_input = torch.randn(2, 4, 8, 8).cpu() torch.onnx.export( net, dummy_input, onnx_file, input_names=['input'], opset_version=opset_version) sess = rt.InferenceSession(onnx_file) onnx_result = sess.run(None, {'input': dummy_input.detach().numpy()}) pytorch_result = func(dummy_input).detach().numpy() assert np.allclose(pytorch_result, onnx_result, atol=1e-3) def test_rotated_feature_align(): if torch.__version__ == 'parrots': pytest.skip('onnx is not supported in parrots directly') try: from mmcv.ops import get_onnxruntime_op_path, rotated_feature_align except (ImportError, ModuleNotFoundError): pytest.skip('rotated_feature_align op is not successfully compiled') ort_custom_op_path = get_onnxruntime_op_path() if not os.path.exists(ort_custom_op_path): pytest.skip('custom ops for onnxruntime are not compiled.') spatial_scale = 1.0 / 8 points = 1 def warpped_function(feature, bbox): return rotated_feature_align( feature, bbox, spatial_scale=spatial_scale, points=points) feature = torch.tensor([[[[1.2924, -0.2172, -0.5222, 0.1172], [0.9144, 1.2248, 1.3115, -0.9690], [-0.8949, -1.1797, -0.9093, -0.3961], [-0.4586, 0.5062, -0.7947, -0.7397]], [[-1.0943, -0.7495, 1.3461, -1.1652], [0.2034, 0.6763, -1.2357, 0.5231], [-1.0062, 1.2592, 1.4225, -0.3951], [-0.1242, -1.6240, 0.1932, 2.7181]], [[-1.6271, -1.0276, 0.0578, -0.2997], [-0.9684, -1.6946, -1.3188, -1.1938], [-1.6744, -0.8917, -0.6556, 1.0073], [-0.1205, 0.3671, -0.3731, -0.5347]]], [[[0.7035, 0.2089, -0.1774, 3.4670], [-0.8505, -0.9278, 1.4714, 0.1644], [0.0898, 0.3531, -0.4007, 0.1927], [1.2569, -0.2636, -0.5223, 0.0616]], [[0.1760, -0.7639, -0.4600, -1.3260], [-0.9921, -0.2970, -0.8955, 1.0508], [1.3515, -0.1641, 1.9679, 1.1986], [-0.3616, 0.6287, 0.4933, 0.3360]], [[-0.5860, 0.2124, -0.8700, 2.4200], [-0.0551, -1.5103, -1.6779, 0.8399], [0.8431, 1.2414, -1.1243, -0.3887], [-2.1254, 0.6047, -0.3515, 0.7254]]]]) bbox = torch.tensor( [[[[1.3080e+01, 1.2688e+01, 1.1214e+01, 9.3944e+01, -9.1905e-01], [3.8104e+01, 1.0134e+01, 1.4659e+02, 9.0306e+01, -9.8211e-01], [-5.3213e+01, 4.9508e+01, 5.1513e+01, 3.2055e+01, -3.1954e-01], [2.6974e+01, 2.5248e+01, 5.4495e+01, 3.1083e+00, -6.2127e-01]], [[-1.5604e+01, -5.1908e+01, 2.3998e+02, 1.5008e+01, -1.2546e+00], [3.1354e+01, -7.3635e+00, 6.7879e+01, 3.5081e+01, -3.3851e-01], [-5.3292e+00, 9.1946e+00, 1.2834e+01, 1.0485e+01, -1.3039e+00], [-2.3925e+01, 3.6623e+01, 3.9875e+01, 7.2009e+01, -6.5934e-01]], [[7.2114e+01, -2.3781e+01, 2.9106e+01, 8.4501e+01, -1.1340e+00], [2.6258e+01, -7.7034e+00, 1.7629e+02, 1.0615e+02, -1.2156e+00], [3.8057e+01, 4.6016e+01, 1.2965e+01, 6.9384e+00, -1.0855e+00], [2.4428e+01, -1.6189e+01, 2.0572e+02, 3.1622e+01, -1.5719e-01]], [[3.8226e+00, 2.9608e+01, 1.4457e+01, 6.8179e+01, -9.1997e-01], [2.5003e+01, -4.2490e+01, 9.6007e+01, 4.9086e+01, -1.4786e+00], [8.5983e+01, 5.4980e+01, 7.8080e+01, 1.0003e+02, -1.0926e+00], [9.9065e+00, 4.1457e+01, 5.9799e+00, 1.7973e+01, -5.6313e-01]]], [[[-1.8244e+01, 4.6309e+00, 5.3010e+01, 2.4310e+01, -7.0345e-01], [1.9419e+01, 3.6704e+01, 5.2390e+01, 5.4133e+01, -3.7730e-01], [5.6387e+01, 2.3752e+01, 9.0441e+00, 1.7792e+01, -1.5583e+00], [3.6303e+01, 1.6396e+01, 2.0283e+01, 1.9148e+01, -8.3419e-01]], [[3.2169e+01, 3.0521e+01, 2.6283e+01, 1.9680e+02, -3.0454e-01], [2.5788e+01, -3.2189e+01, 8.8882e+01, 1.0207e+02, -1.5328e+00], [8.4676e+00, -1.6668e+01, 2.4657e+01, 1.1275e+02, -4.0388e-01], [-1.0799e+01, 6.0422e+00, 9.5807e+00, 3.3677e+01, -3.5438e-01]], [[6.9363e+01, 1.0850e+01, 2.5968e+01, 2.2311e+01, -1.6408e-01], [2.8140e+00, 4.6843e+00, 3.1289e+00, 2.1480e+01, -6.7583e-01], [2.6661e+01, 4.5290e+01, 6.1679e+00, 3.0005e+01, -8.9806e-01], [5.0871e+00, 1.3234e+01, 9.2087e+01, 4.9622e+01, -2.8020e-01]], [[-1.2643e+01, 2.5176e+01, 5.0488e+01, 5.4246e+01, -4.4840e-01], [-3.4521e+01, 9.8435e-01, 5.2413e+01, 9.7996e+00, -8.4218e-01], [4.9829e+01, -1.0808e+01, 2.9848e+01, 7.3579e+01, -6.2672e-01], [8.0446e+01, 2.8064e+01, 4.5273e+01, 5.3809e+01, -1.2359e+00]]]]) # compute pytorch_output with torch.no_grad(): pytorch_output = rotated_feature_align( feature, bbox, spatial_scale=spatial_scale, points=points) # export and load onnx model wrapped_model = WrapFunction(warpped_function) with torch.no_grad(): torch.onnx.export( wrapped_model, (feature, bbox), onnx_file, export_params=True, keep_initializers_as_inputs=True, input_names=['feature', 'bbox'], opset_version=11) onnx_model = onnx.load(onnx_file) session_options = rt.SessionOptions() if os.path.exists(ort_custom_op_path): session_options.register_custom_ops_library(ort_custom_op_path) # 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, session_options) onnx_output = sess.run(None, { 'feature': feature.detach().numpy(), 'bbox': bbox.detach().numpy() }) onnx_output = onnx_output[0] # allclose assert np.allclose(pytorch_output, onnx_output, atol=1e-3) @pytest.mark.parametrize('mode', ['top', 'bottom', 'left', 'right']) def test_corner_pool(mode, opset=11): from mmcv.ops import get_onnxruntime_op_path ort_custom_op_path = get_onnxruntime_op_path() if not os.path.exists(ort_custom_op_path): pytest.skip('custom ops for onnxruntime are not compiled.') from mmcv.ops.corner_pool import CornerPool def corner_pool_func(input): corner_pool_module = CornerPool(mode) return corner_pool_module.corner_pool.apply(input) wrapped_model = WrapFunction(corner_pool_func).eval() input = torch.rand((2, 3, 9, 12)) # (n,c,h,w) with torch.no_grad(): torch.onnx.export( wrapped_model, input, onnx_file, export_params=True, keep_initializers_as_inputs=True, input_names=['input'], output_names=['output'], opset_version=opset) onnx_model = onnx.load(onnx_file) 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) == 1) session_options = rt.SessionOptions() session_options.register_custom_ops_library(ort_custom_op_path) sess = rt.InferenceSession(onnx_file, session_options) ort_result = sess.run(None, {'input': input.detach().numpy()}) pytorch_results = wrapped_model(input.clone()) assert np.allclose(pytorch_results, ort_result, atol=1e-5) @pytest.mark.parametrize('key', ['cummax', 'cummin']) def test_cummax_cummin(key, opset=11): # Note generally `cummax` or `cummin` is exportable to ONNX # as long as the pytorch version >= 1.5.0, since `torch.cummax` # is only supported with torch >= 1.5.0. # But when `cummax` or `cummin` serves as an intermediate component # whose outputs is used as inputs for another modules, it's expected # that pytorch version must be >= 1.7.0. Otherwise error appears like: # `RuntimeError: tuple appears in op that does not forward tuples, # unsupported 'kind: prim::PythonOp`. if version.parse(torch.__version__) < version.parse('1.7.0'): pytest.skip('test_cummax_cummin should be ran with pytorch >= 1.7.0') # register custom op `mmcv::cummax` and `mmcv::cummin` from mmcv.onnx.symbolic import register_extra_symbolics register_extra_symbolics(opset) from mmcv.ops import get_onnxruntime_op_path ort_custom_op_path = get_onnxruntime_op_path() if not os.path.exists(ort_custom_op_path): pytest.skip('custom ops for onnxruntime are not compiled.') input_list = [ # arbitrary shape, e.g. 1-D, 2-D, 3-D, ... torch.rand((2, 3, 4, 1, 5)), torch.rand(1), torch.rand((2, 0, 1)), # tensor.numel() is 0 torch.FloatTensor(), # empty tensor ] cummax_cummin_funcs = {'cummax': torch.cummax, 'cummin': torch.cummin} for input in input_list: ndims = input.dim() # valid dim range is [-ndims, ndims-1] # test for all `dim` value which is valid for dim in range(-ndims, ndims): cummax_func = partial(cummax_cummin_funcs[key], dim=dim) wrapped_model = WrapFunction(cummax_func).eval() with torch.no_grad(): torch.onnx.export( wrapped_model, input, onnx_file, export_params=True, keep_initializers_as_inputs=True, input_names=['input'], output_names=['output', 'indices'], opset_version=opset) onnx_model = onnx.load(onnx_file) 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) == 1) session_options = rt.SessionOptions() session_options.register_custom_ops_library(ort_custom_op_path) sess = rt.InferenceSession(onnx_file, session_options) ort_output, ort_inds = sess.run(None, {'input': input.detach().numpy()}) pytorch_output, pytorch_inds = wrapped_model(input.clone()) pytorch_output = pytorch_output.detach().numpy() pytorch_inds = pytorch_inds.detach().numpy() assert np.allclose(pytorch_output, ort_output, atol=1e-5) assert np.all(pytorch_inds == ort_inds) @pytest.mark.parametrize('shifts_dims_pair', [([-3, 5], [2, 0]), (5, None)]) def test_roll(shifts_dims_pair): opset = 11 from mmcv.onnx.symbolic import register_extra_symbolics register_extra_symbolics(opset) input = torch.arange(0, 4 * 5 * 6, dtype=torch.float32).view(4, 5, 6) shifts, dims = shifts_dims_pair func = partial(torch.roll, shifts=shifts, dims=dims) wrapped_model = WrapFunction(func).eval() with torch.no_grad(): torch.onnx.export( wrapped_model, input, onnx_file, export_params=True, keep_initializers_as_inputs=True, input_names=['input'], output_names=['output'], opset_version=opset) onnx_model = onnx.load(onnx_file) 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) == 1) sess = rt.InferenceSession(onnx_file) ort_output = sess.run(None, {'input': input.detach().numpy()})[0] with torch.no_grad(): pytorch_output = wrapped_model(input.clone()) torch.testing.assert_allclose(ort_output, pytorch_output) @pytest.mark.skipif( not torch.cuda.is_available(), reason='modulated_deform_conv2d only supports in GPU') def test_modulated_deform_conv2d(): try: from mmcv.ops import ModulatedDeformConv2d, get_onnxruntime_op_path except (ImportError, ModuleNotFoundError): pytest.skip('modulated_deform_conv op is not successfully compiled') ort_custom_op_path = get_onnxruntime_op_path() if not os.path.exists(ort_custom_op_path): pytest.skip('custom ops for onnxruntime are not compiled.') # modulated deform conv config in_channels = 3 out_channels = 64 stride = 1 padding = 0 dilation = 1 groups = 1 deform_groups = 1 kernel_size = 3 input = torch.rand(1, in_channels, 28, 28).cuda() # (n, c, h, w) conv_offset = nn.Conv2d( in_channels=3, out_channels=deform_groups * 3 * kernel_size * kernel_size, kernel_size=kernel_size, stride=stride, padding=padding, dilation=dilation, bias=True).cuda() conv_offset.cuda() out = conv_offset(input) o1, o2, mask = torch.chunk(out, 3, dim=1) offset = torch.cat((o1, o2), dim=1) mask = torch.sigmoid(mask) model_with_bias = ModulatedDeformConv2d( in_channels, out_channels, kernel_size, stride, padding, dilation, groups, deform_groups, bias=True) model_without_bias = ModulatedDeformConv2d( in_channels, out_channels, kernel_size, stride, padding, dilation, groups, deform_groups, bias=False) models = [model_with_bias.cuda(), model_without_bias.cuda()] for model in models: # export and load onnx model with torch.no_grad(): torch.onnx.export( model, (input, offset, mask), onnx_file, export_params=True, keep_initializers_as_inputs=True, input_names=['input', 'offset', 'mask'], opset_version=11) session_options = rt.SessionOptions() if os.path.exists(ort_custom_op_path): session_options.register_custom_ops_library(ort_custom_op_path) # compute onnx_output sess = rt.InferenceSession(onnx_file, session_options) onnx_output = sess.run( None, { 'input': input.cpu().detach().numpy(), 'offset': offset.cpu().detach().numpy(), 'mask': mask.cpu().detach().numpy() })[0] # compute pytorch_output with torch.no_grad(): pytorch_output = model(input, offset, mask).cpu() # allclose assert np.allclose(pytorch_output, onnx_output, atol=1e-3) def test_deform_conv2d(threshold=1e-3): try: from mmcv.ops import DeformConv2d, get_onnxruntime_op_path except (ImportError, ModuleNotFoundError): pytest.skip('deform_conv op is not successfully compiled') ort_custom_op_path = get_onnxruntime_op_path() if not os.path.exists(ort_custom_op_path): pytest.skip('custom ops for onnxruntime are not compiled.') # deform conv config # modulated deform conv config in_channels = 1 out_channels = 64 stride = 1 padding = 0 dilation = 1 groups = 1 deform_groups = 1 kernel_size = 2 input = [[[[1., 2., 3.], [0., 1., 2.], [3., 5., 2.]]]] offset_weight = [[[0.1, 0.4, 0.6, 0.1]], [[0.3, 0.2, 0.1, 0.3]], [[0.5, 0.5, 0.2, 0.8]], [[0.8, 0.3, 0.9, 0.1]], [[0.3, 0.1, 0.2, 0.5]], [[0.3, 0.7, 0.5, 0.3]], [[0.6, 0.2, 0.5, 0.3]], [[0.4, 0.1, 0.8, 0.4]]] offset_bias = [0.7, 0.1, 0.8, 0.5, 0.6, 0.5, 0.4, 0.7] deform_weight = [[[0.4, 0.2, 0.1, 0.9]]] x = torch.tensor(input) conv_offset = nn.Conv2d( in_channels=in_channels, out_channels=deform_groups * 2 * kernel_size * kernel_size, kernel_size=kernel_size, stride=stride, padding=padding, dilation=dilation, bias=True) conv_offset.weight.data = torch.nn.Parameter( torch.Tensor(offset_weight).reshape(8, 1, 2, 2)) conv_offset.bias.data = torch.nn.Parameter( torch.Tensor(offset_bias).reshape(8)) offset = conv_offset(x) model = DeformConv2d(in_channels, out_channels, kernel_size, stride, padding, dilation, groups, deform_groups) model.weight.data = torch.nn.Parameter( torch.Tensor(deform_weight).reshape(1, 1, 2, 2)) with torch.no_grad(): torch.onnx.export( model, (x, offset), onnx_file, export_params=True, keep_initializers_as_inputs=True, input_names=['input', 'offset'], opset_version=11) session_options = rt.SessionOptions() if os.path.exists(ort_custom_op_path): session_options.register_custom_ops_library(ort_custom_op_path) # compute onnx_output sess = rt.InferenceSession(onnx_file, session_options) onnx_output = sess.run( None, { 'input': x.cpu().detach().numpy(), 'offset': offset.cpu().detach().numpy(), })[0] # compute pytorch_output with torch.no_grad(): pytorch_output = model(x, offset).cpu() # allclose assert np.allclose(pytorch_output, onnx_output, atol=1e-3)