From b9f2dd4e1025d9a022596a901356f0f8bbc34343 Mon Sep 17 00:00:00 2001 From: Hao Jin Date: Mon, 25 Mar 2019 20:32:23 +0000 Subject: [PATCH] support smooth_l1 in NNVM & Relay --- include/tvm/relay/attrs/nn.h | 10 ++++++ nnvm/include/nnvm/top/nn.h | 9 +++++ nnvm/python/nnvm/frontend/mxnet.py | 2 +- nnvm/python/nnvm/top/nn.py | 12 +++++-- nnvm/src/top/nn/nn.cc | 24 +++++++++++++ .../python/frontend/mxnet/test_forward.py | 10 ++++++ nnvm/tests/python/unittest/test_top_level3.py | 6 ++++ python/tvm/relay/frontend/mxnet.py | 1 + python/tvm/relay/op/nn/_nn.py | 9 +++++ python/tvm/relay/op/nn/nn.py | 28 +++++++++++++++ src/relay/op/nn/nn.cc | 34 +++++++++++++++++++ tests/python/frontend/mxnet/test_forward.py | 8 +++++ tests/python/relay/test_op_level1.py | 26 ++++++++++++++ topi/python/topi/nn/elemwise.py | 24 +++++++++++++ 14 files changed, 200 insertions(+), 3 deletions(-) diff --git a/include/tvm/relay/attrs/nn.h b/include/tvm/relay/attrs/nn.h index 103359e3617c..3c09a76a9c88 100644 --- a/include/tvm/relay/attrs/nn.h +++ b/include/tvm/relay/attrs/nn.h @@ -438,6 +438,16 @@ struct L2NormalizeAttrs : public tvm::AttrsNode { } }; +/*! \brief Attributes for SmoothL1 operator */ +struct SmoothL1Attrs : public tvm::AttrsNode { + double scalar; + + TVM_DECLARE_ATTRS(SmoothL1Attrs, "relay.attrs.SmoothL1Attrs") { + TVM_ATTR_FIELD(scalar).set_default(1.0f) + .describe("scalar input."); + } +}; + } // namespace relay } // namespace tvm #endif // TVM_RELAY_ATTRS_NN_H_ diff --git a/nnvm/include/nnvm/top/nn.h b/nnvm/include/nnvm/top/nn.h index 578f928c5b9f..59f79ede53bd 100644 --- a/nnvm/include/nnvm/top/nn.h +++ b/nnvm/include/nnvm/top/nn.h @@ -505,6 +505,15 @@ struct L2NormalizeParam : public dmlc::Parameter { } }; +struct SmoothL1Param : public dmlc::Parameter { + float scalar; + + DMLC_DECLARE_PARAMETER(SmoothL1Param) { + DMLC_DECLARE_FIELD(scalar).set_default(1.0f) + .describe("scalar input"); + } +}; + } // namespace top } // namespace nnvm diff --git a/nnvm/python/nnvm/frontend/mxnet.py b/nnvm/python/nnvm/frontend/mxnet.py index c9f6777e4898..1ab8da99b308 100644 --- a/nnvm/python/nnvm/frontend/mxnet.py +++ b/nnvm/python/nnvm/frontend/mxnet.py @@ -335,7 +335,7 @@ def _argmin(inputs, attrs): 'flatten', 'log', 'log_softmax', 'max', 'min', 'negative', 'ones_like', 'relu', 'sigmoid', 'slice_like', 'softmax', 'sum', 'tanh', 'transpose', 'zeros_like', 'gather_nd', - 'reshape_like', 'where'] + 'reshape_like', 'where', 'smooth_l1'] _convert_map = { '_copy' : _rename('copy'), diff --git a/nnvm/python/nnvm/top/nn.py b/nnvm/python/nnvm/top/nn.py index 6830db16f04c..210485ee0372 100644 --- a/nnvm/python/nnvm/top/nn.py +++ b/nnvm/python/nnvm/top/nn.py @@ -13,7 +13,6 @@ reg.register_schedule("relu", _fschedule_broadcast) reg.register_pattern("relu", OpPattern.ELEMWISE) - # leaky_relu reg.register_schedule("leaky_relu", _fschedule_broadcast) reg.register_pattern("leaky_relu", OpPattern.ELEMWISE) @@ -36,7 +35,6 @@ reg.register_schedule("__layout_transform__", _fschedule_injective) reg.register_pattern("__layout_transform__", OpPattern.INJECTIVE) - @reg.register_schedule("softmax") def schedule_softmax(_, outs, target): """Schedule definition of softmax""" @@ -426,3 +424,13 @@ def schedule_l2_normalize(attrs, outs, target): return topi.generic.schedule_l2_normalize(outs) reg.register_pattern("l2_normalize", OpPattern.OUT_ELEMWISE_FUSABLE) + +@reg.register_compute("smooth_l1") +def compute_smooth_l1(attrs, inputs, _): + """Compute definition of smooth_l1""" + scalar = attrs.get_float("scalar") + return topi.nn.smooth_l1(inputs[0], scalar) + +# smooth_l1 +reg.register_schedule("smooth_l1", _fschedule_injective) +reg.register_pattern("smooth_l1", OpPattern.INJECTIVE) diff --git a/nnvm/src/top/nn/nn.cc b/nnvm/src/top/nn/nn.cc index f213fa3a19ec..2ca55ecdc1b0 100644 --- a/nnvm/src/top/nn/nn.cc +++ b/nnvm/src/top/nn/nn.cc @@ -779,5 +779,29 @@ NNVM_REGISTER_OP(l2_normalize) .set_attr("FCorrectLayout", ElemwiseArbitraryLayout<1, 1>) .set_support_level(1); +DMLC_REGISTER_PARAMETER(SmoothL1Param); + +inline bool SmoothL1InferShape(const nnvm::NodeAttrs& attrs, + std::vector* in_shape, + std::vector* out_shape) { + TShape dshape = (*in_shape)[0]; + TShape oshape = dshape; + + NNVM_ASSIGN_OUTPUT_SHAPE(attrs, *out_shape, 0, oshape); + return true; +} + +NNVM_REGISTER_OP(smooth_l1) +.describe(R"code(Calculate Smooth L1 Loss)code" NNVM_ADD_FILELINE) +.add_argument("data", "Tensor", "Input data.") +.set_attr_parser(ParamParser) +.set_attr("FGetAttrDict", ParamGetAttrDict) +.set_num_inputs(1) +.set_num_outputs(1) +.set_attr("FInferShape", SmoothL1InferShape) +.set_attr("FInferType", ElemwiseType<1, 1>) +.set_attr("FCorrectLayout", ElemwiseArbitraryLayout<1, 1>) +.set_support_level(1); + } // namespace top } // namespace nnvm diff --git a/nnvm/tests/python/frontend/mxnet/test_forward.py b/nnvm/tests/python/frontend/mxnet/test_forward.py index 581ae75a4bbc..fb170f6dec0c 100644 --- a/nnvm/tests/python/frontend/mxnet/test_forward.py +++ b/nnvm/tests/python/frontend/mxnet/test_forward.py @@ -290,6 +290,15 @@ def test_forward_minimum(): tvm.testing.assert_allclose(mx_out, tvm_out, rtol=1e-5, atol=1e-5) +def test_forward_smooth_l1(): + data = mx.sym.var('data') + mx_sym = mx.sym.smooth_l1(data) + verify_mxnet_frontend_impl(mx_sym, (3, 4), (3, 4)) + data = mx.sym.var('data') + mx_sym = mx.sym.smooth_l1(data, scalar=1.0) + verify_mxnet_frontend_impl(mx_sym, (3, 4), (3, 4)) + + if __name__ == '__main__': test_forward_mlp() test_forward_vgg() @@ -315,3 +324,4 @@ def test_forward_minimum(): test_forward_slice() test_forward_maximum() test_forward_minimum() + test_forward_smooth_l1() diff --git a/nnvm/tests/python/unittest/test_top_level3.py b/nnvm/tests/python/unittest/test_top_level3.py index 47e7a8bce100..05426aca1c12 100644 --- a/nnvm/tests/python/unittest/test_top_level3.py +++ b/nnvm/tests/python/unittest/test_top_level3.py @@ -23,8 +23,14 @@ def test_prelu(): assert(y.list_input_names()[0] == 'x') assert(y.list_input_names()[1] == 'w') +def test_smooth_l1(): + x = sym.Variable("x") + y = sym.smooth_l1(x, scalar=1.0) + assert(y.list_input_names()[0] == 'x') + if __name__ == "__main__": test_scalar_op() test_reshape() test_leaky_relu() test_prelu() + test_smooth_l1() diff --git a/python/tvm/relay/frontend/mxnet.py b/python/tvm/relay/frontend/mxnet.py index 758793c980d6..7804e475ddc7 100644 --- a/python/tvm/relay/frontend/mxnet.py +++ b/python/tvm/relay/frontend/mxnet.py @@ -581,6 +581,7 @@ def _mx_embedding(inputs, _): "zeros_like", "ones_like", "where", + "smooth_l1", ] _convert_map = { diff --git a/python/tvm/relay/op/nn/_nn.py b/python/tvm/relay/op/nn/_nn.py index d38f40ac373b..52138109e614 100644 --- a/python/tvm/relay/op/nn/_nn.py +++ b/python/tvm/relay/op/nn/_nn.py @@ -270,6 +270,15 @@ def schedule_l2_normalize(attrs, outs, target): reg.register_pattern("nn.l2_normalize", OpPattern.OUT_ELEMWISE_FUSABLE) +# smooth_l1 +@reg.register_compute("nn.smooth_l1") +def compute_smooth_l1(attrs, inputs, out_dtype, target): + """Compute definition of smooth_l1""" + return [topi.nn.smooth_l1(inputs[0], attrs.scalar)] + +reg.register_schedule("nn.smooth_l1", schedule_injective) +reg.register_pattern("nn.smooth_l1", OpPattern.INJECTIVE) + # upsampling reg.register_schedule("nn.upsampling", reg.schedule_injective) def schedule_upsampling(_, outs, target): diff --git a/python/tvm/relay/op/nn/nn.py b/python/tvm/relay/op/nn/nn.py index ad8b287bb397..1edbdfbccbb7 100644 --- a/python/tvm/relay/op/nn/nn.py +++ b/python/tvm/relay/op/nn/nn.py @@ -645,6 +645,34 @@ def l2_normalize(data, eps, axis=None): return _make.l2_normalize(data, eps, axis) +def smooth_l1(data, scalar=1.0): +#pylint: disable=anomalous-backslash-in-string + """Compute smooth L1 value on the input data + + .. math:: + y(i, j) = + \begin{cases} + (\sigma x(i, j))^2/2,& \text{if} x(i, j) < 1/\sigma^2\\ + |x(i, j)|-0.5/\sigma^2,& \text{otherwise} + \end{cases} + + Parameters + ---------- + data : tvm.relay.Expr + The input data to the operator. + + scalar : float, default 1.0 + scalar input + + Returns + ------- + result : tvm.relay.Expr + The computed result. + """ +#pylint: enable=anomalous-backslash-in-string + return _make.smooth_l1(data, scalar) + + def dropout(data, rate=0.5): """Applies the dropout operation to the input array. diff --git a/src/relay/op/nn/nn.cc b/src/relay/op/nn/nn.cc index 59f68d9d8880..5c3fa61c1340 100644 --- a/src/relay/op/nn/nn.cc +++ b/src/relay/op/nn/nn.cc @@ -511,6 +511,40 @@ Normalizes along dimension axis using an L2 norm .set_attr("FInferCorrectLayout", ElemwiseArbitraryLayout) .add_type_rel("Identity", IdentityRel); +// smooth_l1 +TVM_REGISTER_NODE_TYPE(SmoothL1Attrs); + +Expr MakeSmoothL1(Expr data, + double scalar) { + auto attrs = make_node(); + attrs->scalar = scalar; + static const Op& op = Op::Get("nn.smooth_l1"); + return CallNode::make(op, {data}, Attrs(attrs), {}); +} + +TVM_REGISTER_API("relay.op.nn._make.smooth_l1") +.set_body([](const TVMArgs& args, TVMRetValue* rv) { + runtime::detail::unpack_call(MakeSmoothL1, args, rv); + }); + +RELAY_REGISTER_OP("nn.smooth_l1") +.describe(R"code(Smooth L1 layer. + +.. math:: + f(x) = + \begin{cases} + (\sigma x)^2/2,& \text{if }x < 1/\sigma^2\\ + |x|-0.5/\sigma^2,& \text{otherwise} + \end{cases} + +)code" TVM_ADD_FILELINE) +.set_attrs_type_key("relay.attrs.SmoothL1Attrs") +.set_num_inputs(1) +.add_argument("data", "Tensor", "The input tensor.") +.set_support_level(2) +.set_attr("FInferCorrectLayout", ElemwiseArbitraryLayout) +.add_type_rel("Identity", IdentityRel); + // Dropout TVM_REGISTER_NODE_TYPE(DropoutAttrs); diff --git a/tests/python/frontend/mxnet/test_forward.py b/tests/python/frontend/mxnet/test_forward.py index aad666ca75b4..9b375df80fae 100644 --- a/tests/python/frontend/mxnet/test_forward.py +++ b/tests/python/frontend/mxnet/test_forward.py @@ -464,6 +464,14 @@ def verify(data_shape, weight_shape): verify((2, 2), (4, 5)) verify((2, 3, 4), (4, 5)) + +def test_forward_smooth_l1(): + data = mx.sym.var('data') + mx_sym = mx.sym.smooth_l1(data) + verify_mxnet_frontend_impl(mx_sym, (3, 4), (3, 4)) + mx_sym = mx.sym.smooth_l1(data, scalar=1.0) + verify_mxnet_frontend_impl(mx_sym, (3, 4), (3, 4)) + if __name__ == '__main__': test_forward_mlp() test_forward_vgg() diff --git a/tests/python/relay/test_op_level1.py b/tests/python/relay/test_op_level1.py index b954e42bf1ab..656826717b56 100644 --- a/tests/python/relay/test_op_level1.py +++ b/tests/python/relay/test_op_level1.py @@ -47,6 +47,32 @@ def check_single_op(opfunc, ref): check_single_op(opfunc, ref) +def test_smooth_l1(): + def smooth_l1(x, scalar=1.0): + res = np.abs(x) - 0.5 / scalar / scalar + res[x < (1.0 / scalar / scalar)] = np.power(scalar * x[x < (1.0 / scalar / scalar)], 2) / 2. + return res + shape = (10, 4) + dtype = 'float32' + tp = relay.TensorType(shape, dtype) + x = relay.var("x", tp) + y = relay.nn.smooth_l1(x) + # test printer + assert ("%0 = {}(%x)".format(y.op.name)) in y.astext() + # test type inference + assert relay.ir_pass.infer_type(y).checked_type == tp + + data = np.random.rand(*shape).astype(dtype) + ref_res = smooth_l1(data) + func = relay.Function([x], y) + for target, ctx in ctx_list(): + # use graph by execuor default for testing, as we need + # create function explicitly to avoid constant-folding. + intrp = relay.create_executor("graph", ctx=ctx, target=target) + op_res = intrp.evaluate(func)(data) + np.testing.assert_allclose(op_res.asnumpy(), ref_res, rtol=1e-3, atol=1e-5) + + def test_binary_op(): def inst(vars, sh): return [vars.get(s, s) for s in sh] diff --git a/topi/python/topi/nn/elemwise.py b/topi/python/topi/nn/elemwise.py index 14a747e67610..c06aaeb4e151 100644 --- a/topi/python/topi/nn/elemwise.py +++ b/topi/python/topi/nn/elemwise.py @@ -77,3 +77,27 @@ def _compute_channelwise(*indices): xval = x(*indices) return tvm.expr.Select(xval > 0, xval, xval * slope(indices[axis])) return tvm.compute(x.shape, _compute_channelwise) + +@tvm.tag_scope(tag=tag.INJECTIVE) +def smooth_l1(x, scalar=1.0): + """ Smooth L1. + It accepts an input ``x`` and computes the output as : + math:`smooth_l1(x) y = x < 1/(\sigma)^2 ? (\sigma x)^2 / 2 : |x| - 0.5 / \sigma^2`, + Arguments: + x : tvm.Tensor + Input argument. + + scalar : float + scalar input. + + Returns: + y : tvm.Tensor + The result. + + """ + def _compute(*indices): + xval = x(*indices) + return tvm.expr.Select(xval < 1.0 / scalar / scalar, + scalar * scalar * xval * xval / 2., + tvm.expr.Select(xval < 0, -xval, xval) - 0.5 / scalar / scalar) + return tvm.compute(x.shape, _compute)