diff --git a/docs/langref/relay_op.rst b/docs/langref/relay_op.rst index d58ba2e66621..e2da42b6ab32 100644 --- a/docs/langref/relay_op.rst +++ b/docs/langref/relay_op.rst @@ -135,6 +135,7 @@ This level enables additional math and transform operators. tvm.relay.vision.multibox_prior tvm.relay.vision.multibox_transform_loc tvm.relay.vision.nms + tvm.relay.vision.yolo_reorg **Level 10: Temporary Operators** @@ -251,6 +252,7 @@ Level 5 Definitions .. autofunction:: tvm.relay.vision.multibox_prior .. autofunction:: tvm.relay.vision.multibox_transform_loc .. autofunction:: tvm.relay.vision.nms +.. autofunction:: tvm.relay.vision.yolo_reorg Level 10 Definitions diff --git a/include/tvm/relay/attrs/vision.h b/include/tvm/relay/attrs/vision.h index d1a5ea41bc69..73b7339e2edb 100644 --- a/include/tvm/relay/attrs/vision.h +++ b/include/tvm/relay/attrs/vision.h @@ -98,6 +98,17 @@ struct ROIAlignAttrs : public tvm::AttrsNode { } }; +/*! \brief Attributes used in yolo reorg operators */ +struct YoloReorgAttrs : public tvm::AttrsNode { + Integer stride; + + TVM_DECLARE_ATTRS(YoloReorgAttrs, "relay.attrs.YoloReorgAttrs") { + TVM_ATTR_FIELD(stride) + .set_default(1) + .describe("Stride value for yolo reorg"); + } +}; + } // namespace relay } // namespace tvm #endif // TVM_RELAY_ATTRS_VISION_H_ diff --git a/python/tvm/relay/op/vision/__init__.py b/python/tvm/relay/op/vision/__init__.py index 710adfeb4955..10cf6c2fd3ee 100644 --- a/python/tvm/relay/op/vision/__init__.py +++ b/python/tvm/relay/op/vision/__init__.py @@ -5,5 +5,7 @@ from .multibox import * from .nms import * from .rcnn import * +from .yolo import * from . import _multibox from . import _rcnn +from . import _yolo diff --git a/python/tvm/relay/op/vision/_yolo.py b/python/tvm/relay/op/vision/_yolo.py new file mode 100644 index 000000000000..749ebfa26dd0 --- /dev/null +++ b/python/tvm/relay/op/vision/_yolo.py @@ -0,0 +1,9 @@ +#pylint: disable=invalid-name, unused-argument +"""Backend compiler related feature registration""" +from __future__ import absolute_import +from ..op import register_schedule, register_pattern +from ..op import schedule_injective, OpPattern + +# reorg +register_pattern("vision.yolo_reorg", OpPattern.INJECTIVE) +register_schedule("vision.yolo_reorg", schedule_injective) diff --git a/python/tvm/relay/op/vision/yolo.py b/python/tvm/relay/op/vision/yolo.py new file mode 100644 index 000000000000..71b7918dca0f --- /dev/null +++ b/python/tvm/relay/op/vision/yolo.py @@ -0,0 +1,34 @@ +"""Yolo operations.""" +from . import _make + +def yolo_reorg(data, stride): + """Yolo reorg operation used in darknet models. + This layer shuffles the input tensor values based on the stride value. + Along with the shuffling, it does the shape transform. + If '(n, c, h, w)' is the data shape and 's' is stride, output shape is '(n, c*s*s, h/s, w/s)' + Example: data(1, 4, 2, 2) = [[[[ 0 1] [ 2 3]] + [[ 4 5] [ 6 7]] + [[ 8 9] [10 11]] + [[12 13] [14 15]]]] + stride = 2 + ret(1, 16, 1, 1) = [[[[ 0]] [[ 2]] [[ 8]] [[10]] + [[ 1]] [[ 3]] [[ 9]] [[11]] + [[ 4]] [[ 6]] [[12]] [[14]] + [[ 5]] [[ 7]] [[13]] [[15]]]] + + Note: stride=1 has no significance for reorg operation. + + Parameters + ---------- + data : relay.Expr + The input data tensor. + + stride : int + The stride value for reorganisation. + + Returns + ------- + ret : relay.Expr + The computed result. + """ + return _make.yolo_reorg(data, stride) diff --git a/src/relay/op/vision/yolo.cc b/src/relay/op/vision/yolo.cc new file mode 100644 index 000000000000..b826d4c6e8e2 --- /dev/null +++ b/src/relay/op/vision/yolo.cc @@ -0,0 +1,78 @@ +/*! + * Copyright (c) 2018 by Contributors + * \file yolo.cc + * \brief Yolo related operators + */ +#include +#include +#include +#include +#include "../op_common.h" +#include "../type_relations.h" + +namespace tvm { +namespace relay { + +TVM_REGISTER_NODE_TYPE(YoloReorgAttrs); + +/*! +* \brief YoloReorgRel Output type and shape relation evaluation function. +* \param num_inputs Number of input types in the args. +* \param attrs The additional attributes of the operator. +* \param reporter The reporter to report solution to. +* \return false if This relation cannot be resolved. true if this relation has been resolved. +*/ +bool YoloReorgRel(const Array& types, + int num_inputs, + const Attrs& attrs, + const TypeReporter& reporter) { + CHECK_EQ(types.size(), 2); + const auto* data = types[0].as(); + if (data == nullptr) return false; + + const YoloReorgAttrs* param = attrs.as(); + CHECK(param != nullptr); + + CHECK(data->shape.size() == 4) << "Yolo reorg supports only 4 dimension."; + std::vector&& oshape = AsVector(data->shape); + oshape[1] = oshape[1] * param->stride * param->stride; + oshape[2] = oshape[2] / param->stride; + oshape[3] = oshape[3] / param->stride; + reporter->Assign(types[1], TensorTypeNode::make(oshape, data->dtype)); + return true; +} + +Expr MakeYoloReorg(Expr data, + Integer stride) { + auto attrs = make_node(); + attrs->stride = stride; + static const Op& op = Op::Get("vision.yolo_reorg"); + return CallNode::make(op, {data}, Attrs(attrs), {}); +} + + +TVM_REGISTER_API("relay.op.vision._make.yolo_reorg") +.set_body([](const TVMArgs& args, TVMRetValue* rv) { + runtime::detail::unpack_call(MakeYoloReorg, args, rv); +}); + + +RELAY_REGISTER_OP("vision.yolo_reorg") +.describe(R"doc("Yolo reorg operation. This layer reorganize the output. +Its function is mostly shape transform.")doc" TVM_ADD_FILELINE) +.add_argument("data", "Tensor", "The input tensor.") +.set_num_inputs(1) +.set_support_level(5) +.set_attrs_type_key("relay.attrs.YoloReorgAttrs") +.add_type_rel("YoloReorg", YoloReorgRel) +.set_attr("FTVMCompute", [](const Attrs& attrs, + const Array& inputs, + const Type& out_type, + const Target& target) { + const auto* params = attrs.as(); + CHECK(params != nullptr); + return Array{ topi::vision::reorg(inputs[0], params->stride) }; +}); + +} // namespace relay +} // namespace tvm diff --git a/tests/python/relay/test_op_level5.py b/tests/python/relay/test_op_level5.py index 1d91d92a6abc..8db6d747ef5e 100644 --- a/tests/python/relay/test_op_level5.py +++ b/tests/python/relay/test_op_level5.py @@ -7,7 +7,6 @@ from tvm.relay.testing import ctx_list import topi.testing - def test_resize_infer_type(): n, c, h, w = tvm.var("n"), tvm.var("c"), tvm.var("h"), tvm.var("w") x = relay.var("x", relay.TensorType((n, c, h, w), "int8")) @@ -307,6 +306,40 @@ def verify_roi_align(data_shape, rois_shape, pooled_size, spatial_scale, sample_ verify_roi_align((4, 4, 16, 16), (32, 5), pooled_size=7, spatial_scale=0.5, sample_ratio=2) +def test_yolo_reorg_infer_shape(): + def verify_yolo_reorg(shape, stride, out_shape): + x = relay.var("x", relay.TensorType(shape, "float32")) + z = relay.vision.yolo_reorg(x, stride=stride) + zz = relay.ir_pass.infer_type(z) + assert "stride=" in z.astext() + assert zz.checked_type == relay.ty.TensorType(out_shape, "float32") + + n, c, h, w = tvm.var("n"), tvm.var("c"), tvm.var("h"), tvm.var("w") + verify_yolo_reorg((n, c, 20, 20), 10, (n, c*10*10, 2, 2)) + verify_yolo_reorg((n, c, h, w), 2, (n, c*2*2, h/2, w/2)) + +def test_yolo_reorg(): + def verify_yolo_reorg(shape, stride): + x_data = np.random.uniform(low=-1, high=1, size=shape).astype("float32") + ref_res = topi.testing.reorg_python(x_data, stride) + + x = relay.var("x", relay.TensorType(shape, "float32")) + z = relay.vision.yolo_reorg(x, stride=stride) + zz = relay.ir_pass.infer_type(z) + assert "stride=" in z.astext() + assert zz.checked_type == relay.ty.TensorType(ref_res.shape, "float32") + + func = relay.Function([x], z) + + for target, ctx in ctx_list(): + for kind in ["graph", "debug"]: + intrp = relay.create_executor(kind, ctx=ctx, target=target) + op_res = intrp.evaluate(func)(x_data) + tvm.testing.assert_allclose(op_res.asnumpy(), ref_res, rtol=1e-5) + + verify_yolo_reorg((1, 100, 20, 20), 10) + verify_yolo_reorg((1, 4, 6, 6), 2) + if __name__ == "__main__": test_resize_infer_type() test_resize() @@ -314,3 +347,5 @@ def verify_roi_align(data_shape, rois_shape, pooled_size, spatial_scale, sample_ test_multibox_transform_loc() test_nms() test_roi_align() + test_yolo_reorg_infer_shape() + test_yolo_reorg()