From 95e732ebefbb6210ae2c68ccc33a01ec2f7ac9a9 Mon Sep 17 00:00:00 2001 From: Wang Date: Thu, 18 Oct 2018 10:43:53 -0700 Subject: [PATCH 1/6] Add Relay NMS operator --- include/tvm/relay/attrs/vision.h | 16 ++++++++ src/relay/op/vision/multibox_op.cc | 1 - src/relay/op/vision/nms.cc | 62 ++++++++++++++++++++++++++++++ 3 files changed, 78 insertions(+), 1 deletion(-) create mode 100644 src/relay/op/vision/nms.cc diff --git a/include/tvm/relay/attrs/vision.h b/include/tvm/relay/attrs/vision.h index 60ee4cb88e43..3ae3d6ad89d0 100644 --- a/include/tvm/relay/attrs/vision.h +++ b/include/tvm/relay/attrs/vision.h @@ -40,6 +40,22 @@ struct MultiBoxPriorAttrs : public tvm::AttrsNode { } }; +/*! \brief Attributes used in non_maximum_suppression operators */ +struct NMSAttrs : public tvm::AttrsNode{ + float nms_threshold; + bool force_suppress; + int nms_topk; + + TVM_DECLARE_ATTRS(NMSAttrs, "relay.attrs.NMSAttrs") { + TVM_ATTR_FIELD(nms_threshold).set_default(0.5) + .describe("Non-maximum suppression threshold."); + TVM_ATTR_FIELD(force_suppress).set_default(false) + .describe("Suppress all detections regardless of class_id."); + TVM_ATTR_FIELD(nms_topk).set_default(-1) + .describe("Keep maximum top k detections before nms, -1 for no limit."); + } +}; + } // namespace relay } // namespace tvm #endif // TVM_RELAY_ATTRS_VISION_H_ diff --git a/src/relay/op/vision/multibox_op.cc b/src/relay/op/vision/multibox_op.cc index ce069a78186b..4a774a4cc161 100644 --- a/src/relay/op/vision/multibox_op.cc +++ b/src/relay/op/vision/multibox_op.cc @@ -5,7 +5,6 @@ */ #include #include -#include namespace tvm { namespace relay { diff --git a/src/relay/op/vision/nms.cc b/src/relay/op/vision/nms.cc new file mode 100644 index 000000000000..40d1def3200c --- /dev/null +++ b/src/relay/op/vision/nms.cc @@ -0,0 +1,62 @@ +/*! + * Copyright (c) 2018 by Contributors + * \file nms.cc + * \brief Non-maximum suppression operators + */ +#include +#include + +namespace tvm { +namespace relay { + +TVM_REGISTER_NODE_TYPE(NMSAttrs); + +bool NMSRel(const Array& types, + int num_inputs, + const Attrs& attrs, + const TypeReporter& reporter) { + CHECK_EQ(types.size(), 3); + const auto* data = types[0].as(); + const auto* valid_count = types[1].as(); + const auto& dshape = data->shape; + const auto& vshape = valid_count->shape; + CHECK_EQ(dshape.size(), 3) << "Input data should be 3-D."; + CHECK_EQ(vshape.size(), 1) << "Input valid count should be 1-D."; + + // assign output type + reporter->Assign(types[2], TensorTypeNode::make(dshape, data->dtype)); + return true; +} + + +Expr MakeNMS(Expr data, + Expr valid_count, + float nms_threshold, + bool force_suppress, + int nms_topk) { + auto attrs = make_node(); + attrs->nms_threshold = std::move(nms_threshold); + attrs->force_suppress= std::move(force_suppress); + attrs->nms_topk = std::move(nms_topk); + static const Op& op = Op::Get("vision.nms"); + return CallNode::make(op, {data, valid_count}, Attrs(attrs), {}); +} + + +TVM_REGISTER_API("relay.op.vision._make.nms") +.set_body([](const TVMArgs& args, TVMRetValue* rv) { + runtime::detail::unpack_call(MakeNMS, args, rv); +}); + + +RELAY_REGISTER_OP("vision.nms") +.describe(R"doc("Non-maximum suppression." +)doc" TVM_ADD_FILELINE) +.set_num_inputs(2) +.add_argument("data", "Tensor", "Input data.") +.add_argument("valid_count", "Tensor", "Number of valid anchor boxes.") +.set_support_level(4) +.add_type_rel("NMS", NMSRel); + +} // namespace relay +} // namespace tvm \ No newline at end of file From 8ac211ce7d523365628e3553dadd05fceb93a537 Mon Sep 17 00:00:00 2001 From: Wang Date: Thu, 18 Oct 2018 16:57:13 -0700 Subject: [PATCH 2/6] Add tests --- include/tvm/relay/attrs/vision.h | 2 +- python/tvm/relay/op/vision/__init__.py | 1 + python/tvm/relay/op/vision/nms.py | 36 ++++++++++++++++++++++++++ src/relay/op/vision/multibox_op.cc | 2 +- src/relay/op/vision/nms.cc | 14 +++++----- tests/python/relay/test_op_level5.py | 36 +++++++++++++++++++++++++- 6 files changed, 81 insertions(+), 10 deletions(-) create mode 100644 python/tvm/relay/op/vision/nms.py diff --git a/include/tvm/relay/attrs/vision.h b/include/tvm/relay/attrs/vision.h index 3ae3d6ad89d0..d69785b976fd 100644 --- a/include/tvm/relay/attrs/vision.h +++ b/include/tvm/relay/attrs/vision.h @@ -42,7 +42,7 @@ struct MultiBoxPriorAttrs : public tvm::AttrsNode { /*! \brief Attributes used in non_maximum_suppression operators */ struct NMSAttrs : public tvm::AttrsNode{ - float nms_threshold; + double nms_threshold; bool force_suppress; int nms_topk; diff --git a/python/tvm/relay/op/vision/__init__.py b/python/tvm/relay/op/vision/__init__.py index b3010d2d5310..9ecd8a84770a 100644 --- a/python/tvm/relay/op/vision/__init__.py +++ b/python/tvm/relay/op/vision/__init__.py @@ -3,3 +3,4 @@ from __future__ import absolute_import as _abs from .multibox import * +from .nms import * diff --git a/python/tvm/relay/op/vision/nms.py b/python/tvm/relay/op/vision/nms.py new file mode 100644 index 000000000000..5cbcaff4c5e6 --- /dev/null +++ b/python/tvm/relay/op/vision/nms.py @@ -0,0 +1,36 @@ +"""Non-maximum suppression operations.""" +from __future__ import absolute_import as _abs +from . import _make + +def nms(data, + valid_count, + nms_threshold=0.5, + force_suppress=False, + nms_topk=-1): + """Generate prior(anchor) boxes from data, sizes and ratios. + + Parameters + ---------- + data : relay.Expr + 3-D tensor with shape [batch_size, num_anchors, 6]. + The last dimension should be in format of + [class_id, score, box_left, box_top, box_right, box_bottom]. + + valid_count : relay.Expr + Number of valid anchor boxes. + + nms_threshold : float, optional + Non-maximum suppression threshold. + + force_suppress : bool, optional + Suppress all detections regardless of class_id. + + nms_topk : int, optional + Keep maximum top k detections before nms, -1 for no limit. + + Returns + ------- + out : relay.Expr + 3-D tensor with shape [batch_size, num_anchors, 6]. + """ + return _make.nms(data, valid_count, nms_threshold, force_suppress, nms_topk) diff --git a/src/relay/op/vision/multibox_op.cc b/src/relay/op/vision/multibox_op.cc index 4a774a4cc161..e347e544e4f9 100644 --- a/src/relay/op/vision/multibox_op.cc +++ b/src/relay/op/vision/multibox_op.cc @@ -65,7 +65,7 @@ RELAY_REGISTER_OP("vision.multibox_prior") .set_attrs_type_key("relay.attrs.MultiBoxPriorAttrs") .set_num_inputs(1) .add_argument("data", "Tensor", "The input tensor.") -.set_support_level(4) +.set_support_level(5) .add_type_rel("MultiBoxPrior", MultiboxPriorRel); } // namespace relay diff --git a/src/relay/op/vision/nms.cc b/src/relay/op/vision/nms.cc index 40d1def3200c..3f1f911906b5 100644 --- a/src/relay/op/vision/nms.cc +++ b/src/relay/op/vision/nms.cc @@ -31,13 +31,13 @@ bool NMSRel(const Array& types, Expr MakeNMS(Expr data, Expr valid_count, - float nms_threshold, + double nms_threshold, bool force_suppress, int nms_topk) { auto attrs = make_node(); - attrs->nms_threshold = std::move(nms_threshold); - attrs->force_suppress= std::move(force_suppress); - attrs->nms_topk = std::move(nms_topk); + attrs->nms_threshold = nms_threshold; + attrs->force_suppress = force_suppress; + attrs->nms_topk = nms_topk; static const Op& op = Op::Get("vision.nms"); return CallNode::make(op, {data, valid_count}, Attrs(attrs), {}); } @@ -45,7 +45,7 @@ Expr MakeNMS(Expr data, TVM_REGISTER_API("relay.op.vision._make.nms") .set_body([](const TVMArgs& args, TVMRetValue* rv) { - runtime::detail::unpack_call(MakeNMS, args, rv); + runtime::detail::unpack_call(MakeNMS, args, rv); }); @@ -55,8 +55,8 @@ RELAY_REGISTER_OP("vision.nms") .set_num_inputs(2) .add_argument("data", "Tensor", "Input data.") .add_argument("valid_count", "Tensor", "Number of valid anchor boxes.") -.set_support_level(4) +.set_support_level(5) .add_type_rel("NMS", NMSRel); } // namespace relay -} // namespace tvm \ No newline at end of file +} // namespace tvm diff --git a/tests/python/relay/test_op_level5.py b/tests/python/relay/test_op_level5.py index 4e554cd0cf81..e51ab516e55d 100644 --- a/tests/python/relay/test_op_level5.py +++ b/tests/python/relay/test_op_level5.py @@ -18,7 +18,6 @@ def test_resize_infer_type(): assert zz.checked_type == relay.TensorType((n, c, 100, 200), "int8") - def test_multibox_prior(): sizes = (0.3, 1.5, 0.7) ratios = (1.3, 2.4) @@ -44,6 +43,41 @@ def test_multibox_prior(): (1, h * w, 4), "float32") +def test_nms(): + num_anchors = 60 + + nms_threshold = 0.5 + force_suppress = True + nms_topk = 10 + + ib = relay.ir_builder.IRBuilder() + n = tvm.var("n") + x0 = ib.param("x0", relay.ty.TensorType((n, num_anchors, 6), "float32")) + x1 = ib.param("x1", relay.ty.TensorType((n,), "int")) + + with ib.function(x0, x1) as func: + ib.ret(relay.vision.nms(x0, x1, nms_threshold, force_suppress, nms_topk)) + ib.ret(func) + func = relay.ir_pass.infer_type(ib.env, func.to_func()) + ftype = func.checked_type + assert ftype.ret_type == relay.ty.TensorType( + (n, num_anchors, 6), "float32") + + ib = relay.ir_builder.IRBuilder() + n = tvm.var("n") + x0 = ib.param("x0", relay.ty.TensorType((n, num_anchors, 6), "float32")) + x1 = ib.param("x1", relay.ty.TensorType((n,), "int")) + + with ib.function(x0, x1) as func: + ib.ret(relay.vision.nms(x0, x1)) + ib.ret(func) + func = relay.ir_pass.infer_type(ib.env, func.to_func()) + ftype = func.checked_type + assert ftype.ret_type == relay.ty.TensorType( + (n, num_anchors, 6), "float32") + + if __name__ == "__main__": test_resize_infer_type() test_multibox_prior() + test_nms() From 9ed7dd48ea611ab989df89b28111581ea0bac544 Mon Sep 17 00:00:00 2001 From: Wang Date: Wed, 24 Oct 2018 12:51:38 -0700 Subject: [PATCH 3/6] Rename threshold arg --- include/tvm/relay/attrs/vision.h | 4 ++-- python/tvm/relay/op/vision/nms.py | 8 ++++---- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/include/tvm/relay/attrs/vision.h b/include/tvm/relay/attrs/vision.h index d69785b976fd..d3e35be430cb 100644 --- a/include/tvm/relay/attrs/vision.h +++ b/include/tvm/relay/attrs/vision.h @@ -42,12 +42,12 @@ struct MultiBoxPriorAttrs : public tvm::AttrsNode { /*! \brief Attributes used in non_maximum_suppression operators */ struct NMSAttrs : public tvm::AttrsNode{ - double nms_threshold; + double overlap_threshold; bool force_suppress; int nms_topk; TVM_DECLARE_ATTRS(NMSAttrs, "relay.attrs.NMSAttrs") { - TVM_ATTR_FIELD(nms_threshold).set_default(0.5) + TVM_ATTR_FIELD(overlap_threshold).set_default(0.5) .describe("Non-maximum suppression threshold."); TVM_ATTR_FIELD(force_suppress).set_default(false) .describe("Suppress all detections regardless of class_id."); diff --git a/python/tvm/relay/op/vision/nms.py b/python/tvm/relay/op/vision/nms.py index 5cbcaff4c5e6..24a4d638fc48 100644 --- a/python/tvm/relay/op/vision/nms.py +++ b/python/tvm/relay/op/vision/nms.py @@ -4,10 +4,10 @@ def nms(data, valid_count, - nms_threshold=0.5, + overlap_threshold=0.5, force_suppress=False, nms_topk=-1): - """Generate prior(anchor) boxes from data, sizes and ratios. + """Non-maximum suppression operator for object detection. Parameters ---------- @@ -17,7 +17,7 @@ def nms(data, [class_id, score, box_left, box_top, box_right, box_bottom]. valid_count : relay.Expr - Number of valid anchor boxes. + 1-D tensor for valid number of boxes. nms_threshold : float, optional Non-maximum suppression threshold. @@ -33,4 +33,4 @@ def nms(data, out : relay.Expr 3-D tensor with shape [batch_size, num_anchors, 6]. """ - return _make.nms(data, valid_count, nms_threshold, force_suppress, nms_topk) + return _make.nms(data, valid_count, overlap_threshold, force_suppress, nms_topk) From 244913b7f7ed87f8d53c883bd7ddf48c381e8506 Mon Sep 17 00:00:00 2001 From: Wang Date: Wed, 24 Oct 2018 13:12:29 -0700 Subject: [PATCH 4/6] Fix build --- src/relay/op/vision/nms.cc | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/relay/op/vision/nms.cc b/src/relay/op/vision/nms.cc index 3f1f911906b5..a57c6a39f909 100644 --- a/src/relay/op/vision/nms.cc +++ b/src/relay/op/vision/nms.cc @@ -31,11 +31,11 @@ bool NMSRel(const Array& types, Expr MakeNMS(Expr data, Expr valid_count, - double nms_threshold, + double overlap_threshold, bool force_suppress, int nms_topk) { auto attrs = make_node(); - attrs->nms_threshold = nms_threshold; + attrs->overlap_threshold = overlap_threshold; attrs->force_suppress = force_suppress; attrs->nms_topk = nms_topk; static const Op& op = Op::Get("vision.nms"); From 317302fb94a2f7f053e1d2474c5e5faddfadc445 Mon Sep 17 00:00:00 2001 From: Wang Date: Wed, 24 Oct 2018 13:47:21 -0700 Subject: [PATCH 5/6] Fix test --- tests/python/relay/test_op_level5.py | 38 +++++++++++++--------------- 1 file changed, 17 insertions(+), 21 deletions(-) diff --git a/tests/python/relay/test_op_level5.py b/tests/python/relay/test_op_level5.py index e51ab516e55d..c1f6e0212f6f 100644 --- a/tests/python/relay/test_op_level5.py +++ b/tests/python/relay/test_op_level5.py @@ -46,34 +46,30 @@ def test_multibox_prior(): def test_nms(): num_anchors = 60 - nms_threshold = 0.5 + overlap_threshold = 0.5 force_suppress = True nms_topk = 10 - ib = relay.ir_builder.IRBuilder() n = tvm.var("n") - x0 = ib.param("x0", relay.ty.TensorType((n, num_anchors, 6), "float32")) - x1 = ib.param("x1", relay.ty.TensorType((n,), "int")) - - with ib.function(x0, x1) as func: - ib.ret(relay.vision.nms(x0, x1, nms_threshold, force_suppress, nms_topk)) - ib.ret(func) - func = relay.ir_pass.infer_type(ib.env, func.to_func()) - ftype = func.checked_type - assert ftype.ret_type == relay.ty.TensorType( + x0 = relay.var("x0", relay.ty.TensorType((n, num_anchors, 6), "float32")) + x1 = relay.var("x1", relay.ty.TensorType((n,), "int")) + + z = relay.vision.nms(x0, x1, overlap_threshold, force_suppress, nms_topk) + + print(z.astext()) + assert "overlap_threshold" in z.astext() + zz = relay.ir_pass.infer_type(z) + assert zz.checked_type == relay.ty.TensorType( (n, num_anchors, 6), "float32") - ib = relay.ir_builder.IRBuilder() n = tvm.var("n") - x0 = ib.param("x0", relay.ty.TensorType((n, num_anchors, 6), "float32")) - x1 = ib.param("x1", relay.ty.TensorType((n,), "int")) - - with ib.function(x0, x1) as func: - ib.ret(relay.vision.nms(x0, x1)) - ib.ret(func) - func = relay.ir_pass.infer_type(ib.env, func.to_func()) - ftype = func.checked_type - assert ftype.ret_type == relay.ty.TensorType( + x0 = relay.var("x0", relay.ty.TensorType((n, num_anchors, 6), "float32")) + x1 = relay.var("x1", relay.ty.TensorType((n,), "int")) + + z = relay.vision.nms(x0, x1) + + zz = relay.ir_pass.infer_type(z) + assert zz.checked_type == relay.ty.TensorType( (n, num_anchors, 6), "float32") From 8dbdaec560579acc0abe775c7037882c714718a5 Mon Sep 17 00:00:00 2001 From: Wang Date: Tue, 13 Nov 2018 11:27:26 -0800 Subject: [PATCH 6/6] Rename argument --- include/tvm/relay/attrs/vision.h | 4 ++-- python/tvm/relay/op/vision/nms.py | 8 ++++---- src/relay/op/vision/nms.cc | 4 ++-- tests/python/relay/test_op_level5.py | 1 - 4 files changed, 8 insertions(+), 9 deletions(-) diff --git a/include/tvm/relay/attrs/vision.h b/include/tvm/relay/attrs/vision.h index d3e35be430cb..5408582c8356 100644 --- a/include/tvm/relay/attrs/vision.h +++ b/include/tvm/relay/attrs/vision.h @@ -44,14 +44,14 @@ struct MultiBoxPriorAttrs : public tvm::AttrsNode { struct NMSAttrs : public tvm::AttrsNode{ double overlap_threshold; bool force_suppress; - int nms_topk; + int topk; TVM_DECLARE_ATTRS(NMSAttrs, "relay.attrs.NMSAttrs") { TVM_ATTR_FIELD(overlap_threshold).set_default(0.5) .describe("Non-maximum suppression threshold."); TVM_ATTR_FIELD(force_suppress).set_default(false) .describe("Suppress all detections regardless of class_id."); - TVM_ATTR_FIELD(nms_topk).set_default(-1) + TVM_ATTR_FIELD(topk).set_default(-1) .describe("Keep maximum top k detections before nms, -1 for no limit."); } }; diff --git a/python/tvm/relay/op/vision/nms.py b/python/tvm/relay/op/vision/nms.py index 24a4d638fc48..8035e3030b17 100644 --- a/python/tvm/relay/op/vision/nms.py +++ b/python/tvm/relay/op/vision/nms.py @@ -6,7 +6,7 @@ def nms(data, valid_count, overlap_threshold=0.5, force_suppress=False, - nms_topk=-1): + topk=-1): """Non-maximum suppression operator for object detection. Parameters @@ -19,13 +19,13 @@ def nms(data, valid_count : relay.Expr 1-D tensor for valid number of boxes. - nms_threshold : float, optional + overlap_threshold : float, optional Non-maximum suppression threshold. force_suppress : bool, optional Suppress all detections regardless of class_id. - nms_topk : int, optional + topk : int, optional Keep maximum top k detections before nms, -1 for no limit. Returns @@ -33,4 +33,4 @@ def nms(data, out : relay.Expr 3-D tensor with shape [batch_size, num_anchors, 6]. """ - return _make.nms(data, valid_count, overlap_threshold, force_suppress, nms_topk) + return _make.nms(data, valid_count, overlap_threshold, force_suppress, topk) diff --git a/src/relay/op/vision/nms.cc b/src/relay/op/vision/nms.cc index a57c6a39f909..3e3f73bc6cb4 100644 --- a/src/relay/op/vision/nms.cc +++ b/src/relay/op/vision/nms.cc @@ -33,11 +33,11 @@ Expr MakeNMS(Expr data, Expr valid_count, double overlap_threshold, bool force_suppress, - int nms_topk) { + int topk) { auto attrs = make_node(); attrs->overlap_threshold = overlap_threshold; attrs->force_suppress = force_suppress; - attrs->nms_topk = nms_topk; + attrs->topk = topk; static const Op& op = Op::Get("vision.nms"); return CallNode::make(op, {data, valid_count}, Attrs(attrs), {}); } diff --git a/tests/python/relay/test_op_level5.py b/tests/python/relay/test_op_level5.py index c1f6e0212f6f..0bd7a4816a1b 100644 --- a/tests/python/relay/test_op_level5.py +++ b/tests/python/relay/test_op_level5.py @@ -56,7 +56,6 @@ def test_nms(): z = relay.vision.nms(x0, x1, overlap_threshold, force_suppress, nms_topk) - print(z.astext()) assert "overlap_threshold" in z.astext() zz = relay.ir_pass.infer_type(z) assert zz.checked_type == relay.ty.TensorType(