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()
})
os.remove(onnx_file)
assert np.allclose(pytorch_dets, onnx_dets, atol=1e-3)
assert np.allclose(onnx_inds, onnx_inds, atol=1e-3)
def test_softnms_allclose(self):
if not torch.cuda.is_available():
return
from mmcv.ops import soft_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)
np_output = {
'linear': {
'dets':
np.array(
[[3., 6., 9., 11., 0.9], [6., 3., 8., 7., 0.6],
[3., 7., 10., 12., 0.29024392], [1., 4., 13., 7., 0.2]],
dtype=np.float32),
'inds':
np.array([1, 0, 2, 3], dtype=np.int64)
},
'gaussian': {
'dets':
np.array([[3., 6., 9., 11., 0.9], [6., 3., 8., 7., 0.59630775],
[3., 7., 10., 12., 0.35275510],
[1., 4., 13., 7., 0.18650459]],
dtype=np.float32),
'inds':
np.array([1, 0, 2, 3], dtype=np.int64)
},
'naive': {
'dets':
np.array([[3., 6., 9., 11., 0.9], [6., 3., 8., 7., 0.6],
[1., 4., 13., 7., 0.2]],
dtype=np.float32),
'inds':
np.array([1, 0, 3], dtype=np.int64)
}
}
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']]
for iou, sig, mscore, m in configs:
dets, inds = soft_nms(boxes,
scores,
iou_threshold=iou,
sigma=sig,
min_score=mscore,
method=m)
assert np.allclose(dets.cpu().numpy(), np_output[m]['dets'])
assert np.allclose(inds.cpu().numpy(), np_output[m]['inds'])
if torch.__version__ != 'parrots':
boxes = boxes.cuda()
scores = scores.cuda()
for iou, sig, mscore, m in configs:
dets, inds = soft_nms(boxes,
scores,
iou_threshold=iou,
sigma=sig,
min_score=mscore,
method=m)
assert np.allclose(dets.cpu().numpy(), np_output[m]['dets'])
assert np.allclose(inds.cpu().numpy(), np_output[m]['inds'])
