Support MXNet smooth_l1

include/tvm/relay/attrs/nn.h

@@ -438,6 +438,16 @@ struct L2NormalizeAttrs : public tvm::AttrsNode<L2NormalizeAttrs> {
   }
 };
 
+/*! \brief Attributes for SmoothL1 operator */
+struct SmoothL1Attrs : public tvm::AttrsNode<SmoothL1Attrs> {
+  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_

nnvm/include/nnvm/top/nn.h

@@ -505,6 +505,15 @@ struct L2NormalizeParam : public dmlc::Parameter<L2NormalizeParam> {
   }
 };
 
+struct SmoothL1Param : public dmlc::Parameter<SmoothL1Param> {
+  float scalar;
+
+  DMLC_DECLARE_PARAMETER(SmoothL1Param) {
+    DMLC_DECLARE_FIELD(scalar).set_default(1.0f)
+      .describe("scalar input");
+  }
+};
+
 }  // namespace top
 }  // namespace nnvm
 

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'),

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)

nnvm/src/top/nn/nn.cc

@@ -779,5 +779,29 @@ NNVM_REGISTER_OP(l2_normalize)
 .set_attr<FCorrectLayout>("FCorrectLayout", ElemwiseArbitraryLayout<1, 1>)
 .set_support_level(1);
 
+DMLC_REGISTER_PARAMETER(SmoothL1Param);
+
+inline bool SmoothL1InferShape(const nnvm::NodeAttrs& attrs,
+                               std::vector<TShape>* in_shape,
+                               std::vector<TShape>* 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<SmoothL1Param>)
+.set_attr<FGetAttrDict>("FGetAttrDict", ParamGetAttrDict<SmoothL1Param>)
+.set_num_inputs(1)
+.set_num_outputs(1)
+.set_attr<FInferShape>("FInferShape", SmoothL1InferShape)
+.set_attr<FInferType>("FInferType", ElemwiseType<1, 1>)
+.set_attr<FCorrectLayout>("FCorrectLayout", ElemwiseArbitraryLayout<1, 1>)
+.set_support_level(1);
+
 }  // namespace top
 }  // namespace nnvm

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()

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()

python/tvm/relay/frontend/mxnet.py

@@ -581,6 +581,7 @@ def _mx_embedding(inputs, _):
     "zeros_like",
     "ones_like",
     "where",
+    "smooth_l1",
 ]
 
 _convert_map = {

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):

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.
 

src/relay/op/nn/nn.cc

@@ -511,6 +511,40 @@ Normalizes along dimension axis using an L2 norm
 .set_attr<FInferCorrectLayout>("FInferCorrectLayout", ElemwiseArbitraryLayout)
 .add_type_rel("Identity", IdentityRel);
 
+// smooth_l1
+TVM_REGISTER_NODE_TYPE(SmoothL1Attrs);
+
+Expr MakeSmoothL1(Expr data,
+                  double scalar) {
+  auto attrs = make_node<SmoothL1Attrs>();
+  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<Expr, 2>(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>("FInferCorrectLayout", ElemwiseArbitraryLayout)
+.add_type_rel("Identity", IdentityRel);
+
 // Dropout
 TVM_REGISTER_NODE_TYPE(DropoutAttrs);
 

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()

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]

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)