diff --git a/include/tvm/relay/attrs/transform.h b/include/tvm/relay/attrs/transform.h index cc97a94a1406..44d9b1a8ebec 100644 --- a/include/tvm/relay/attrs/transform.h +++ b/include/tvm/relay/attrs/transform.h @@ -478,6 +478,16 @@ struct UniqueAttrs : public tvm::AttrsNode { } }; // struct UniqueAttrs +/*! \brief Attributes used in segment_max, segment_min, + segment_mean, segment_sum, segment_prod operator */ +struct SegmentAttrs : public tvm::AttrsNode { + int num_segments; + + TVM_DECLARE_ATTRS(SegmentAttrs, "relay.attrs.SegmentAttrs") { + TVM_ATTR_FIELD(num_segments).set_default(0).describe("The maximum of segment_ids."); + } +}; // struct SegmentAttrs + } // namespace relay } // namespace tvm #endif // TVM_RELAY_ATTRS_TRANSFORM_H_ diff --git a/python/tvm/relay/frontend/tensorflow.py b/python/tvm/relay/frontend/tensorflow.py index 4af73702ad9c..9f2dde317166 100644 --- a/python/tvm/relay/frontend/tensorflow.py +++ b/python/tvm/relay/frontend/tensorflow.py @@ -912,6 +912,22 @@ def _impl(inputs, attr, params, mod): return _impl +def _unsorted_segment(name): + def _impl(inputs, attr, params, mod): + # op description: https://www.tensorflow.org/api_docs/python/tf/math/unsorted_segment_max + try: + num_segments = _infer_value(inputs[2], params).asnumpy().tolist() + except Exception: + raise tvm.error.OpAttributeInvalid("Can't find num_segments.") + return AttrCvt( + op_name="segment_" + name, + ignores=["Tdim", "Tidx", "Tindices", "Tnumsegments"], + extras={"num_segments": num_segments}, + )([inputs[0], inputs[1]], attr) + + return _impl + + def _crop_and_resize(): def _impl(inputs, attr, params, mod): # input image is a 4-D tensor of shape [batch, image_height, image_width, depth] @@ -2617,6 +2633,23 @@ def _impl(inputs, attr, params, mod): return _impl +def _segment(opname): + def _impl(inputs, attr, params, mod): + # op description: https://www.tensorflow.org/api_docs/python/tf/math/segment_max + try: + segment_ids = _infer_value(inputs[1], params) + except Exception: + raise tvm.error.OpAttributeInvalid("Can't get value of segment_ids.") + + num_out = segment_ids.asnumpy().max() + 1 + out = AttrCvt(op_name=opname, ignores=["T", "Tindices"], extras={"num_segments": num_out})( + inputs, attr + ) + return out + + return _impl + + def _size(): def _impl(inputs, attr, params, mod): new_attr = attr @@ -2864,6 +2897,11 @@ def _impl(inputs, attr, params, mod): "SelectV2": _where(), "Selu": _selu(), "Shape": _shape(), + "SegmentMax": _segment("segment_max"), + "SegmentMean": _segment("segment_mean"), + "SegmentMin": _segment("segment_min"), + "SegmentProd": _segment("segment_prod"), + "SegmentSum": _segment("segment_sum"), "Sigmoid": AttrCvt("sigmoid"), "Sign": AttrCvt("sign"), "Sin": AttrCvt("sin"), @@ -2915,6 +2953,11 @@ def _impl(inputs, attr, params, mod): "UniqueWithCounts": _unique(True), "Unpack": _unpack(), "UnravelIndex": _unravel_index(), + "UnsortedSegmentMax": _unsorted_segment("max"), + "UnsortedSegmentMin": _unsorted_segment("min"), + "UnsortedSegmentMean": _unsorted_segment("mean"), + "UnsortedSegmentProd": _unsorted_segment("prod"), + "UnsortedSegmentSum": _unsorted_segment("sum"), "Where": _where(), "ZerosLike": AttrCvt("zeros_like"), } diff --git a/python/tvm/relay/op/_tensor.py b/python/tvm/relay/op/_tensor.py index d7d99c017b2b..88869abf9e22 100644 --- a/python/tvm/relay/op/_tensor.py +++ b/python/tvm/relay/op/_tensor.py @@ -21,7 +21,8 @@ from tvm import topi from tvm.runtime import convert -from .op import register_compute, register_shape_func +from . import strategy +from .op import register_compute, register_shape_func, register_strategy from .op import register_broadcast_schedule, register_injective_schedule from .op import register_pattern, OpPattern @@ -283,3 +284,23 @@ def elemwise_shape_func(attrs, inputs, _): register_shape_func("sigmoid", False, elemwise_shape_func) register_shape_func("tanh", False, elemwise_shape_func) register_shape_func("logical_not", False, elemwise_shape_func) + +# segment_max +register_strategy("segment_max", strategy.segment_max_strategy) +register_pattern("segment_max", OpPattern.OPAQUE) + +# segment_min +register_strategy("segment_min", strategy.segment_min_strategy) +register_pattern("segment_min", OpPattern.OPAQUE) + +# segment_mean +register_strategy("segment_mean", strategy.segment_mean_strategy) +register_pattern("segment_mean", OpPattern.OPAQUE) + +# segment_sum +register_strategy("segment_sum", strategy.segment_sum_strategy) +register_pattern("segment_sum", OpPattern.OPAQUE) + +# segment_prod +register_strategy("segment_prod", strategy.segment_prod_strategy) +register_pattern("segment_prod", OpPattern.OPAQUE) diff --git a/python/tvm/relay/op/op_attrs.py b/python/tvm/relay/op/op_attrs.py index 6844d133a77e..8deafd6af598 100644 --- a/python/tvm/relay/op/op_attrs.py +++ b/python/tvm/relay/op/op_attrs.py @@ -329,6 +329,11 @@ class ProposalAttrs(Attrs): """Attributes used in proposal operators""" +@tvm._ffi.register_object("relay.attrs.SegmentAttrs") +class SegmentAttrs(Attrs): + """Attributes used in segment operators""" + + @tvm._ffi.register_object("relay.attrs.MaxPool2DAttrs") class MaxPool2DAttrs(Attrs): """Attributes used in max_pool2d operators""" diff --git a/python/tvm/relay/op/strategy/generic.py b/python/tvm/relay/op/strategy/generic.py index a6ad06e544a6..8ea5c95752e5 100644 --- a/python/tvm/relay/op/strategy/generic.py +++ b/python/tvm/relay/op/strategy/generic.py @@ -1540,6 +1540,29 @@ def uniform_strategy(attrs, inputs, out_type, target): return strategy +# segment_max +def wrap_compute_segment_max(topi_compute): + """wrap segment_max topi compute""" + + def _compute_segment_max(attrs, inputs, out_type): + num_segments = attrs.num_segments + return [topi_compute(inputs[0], inputs[1], num_segments, "max")] + + return _compute_segment_max + + +@override_native_generic_func("segment_max_strategy") +def segment_max_strategy(attrs, inputs, out_type, target): + """segment_max generic strategy""" + strategy = _op.OpStrategy() + strategy.add_implementation( + wrap_compute_segment_max(topi.segment_op), + wrap_topi_schedule(topi.generic.schedule_segment_op), + name="segment_max.generic", + ) + return strategy + + def wrap_compute_scanop(topi_compute): """Wrap scanop style topi compute""" @@ -1561,6 +1584,29 @@ def cumsum_strategy(attrs, inputs, out_type, target): return strategy +# segment_min +def wrap_compute_segment_min(topi_compute): + """wrap segment_min topi compute""" + + def _compute_segment_min(attrs, inputs, out_type): + num_segments = attrs.num_segments + return [topi_compute(inputs[0], inputs[1], num_segments, "min")] + + return _compute_segment_min + + +@override_native_generic_func("segment_min_strategy") +def segment_min_strategy(attrs, inputs, out_type, target): + """segment_min generic strategy""" + strategy = _op.OpStrategy() + strategy.add_implementation( + wrap_compute_segment_min(topi.segment_op), + wrap_topi_schedule(topi.generic.schedule_segment_op), + name="segment_min.generic", + ) + return strategy + + @override_native_generic_func("cumprod_strategy") def cumprod_strategy(attrs, inputs, out_type, target): """cumprod generic strategy""" @@ -1573,6 +1619,29 @@ def cumprod_strategy(attrs, inputs, out_type, target): return strategy +# segment_mean +def wrap_compute_segment_mean(topi_compute): + """wrap segment_mean topi compute""" + + def _compute_segment_mean(attrs, inputs, out_type): + num_segments = attrs.num_segments + return [topi_compute(inputs[0], inputs[1], num_segments, "mean")] + + return _compute_segment_mean + + +@override_native_generic_func("segment_mean_strategy") +def segment_mean_strategy(attrs, inputs, out_type, target): + """segment_mean generic strategy""" + strategy = _op.OpStrategy() + strategy.add_implementation( + wrap_compute_segment_mean(topi.segment_op), + wrap_topi_schedule(topi.generic.schedule_segment_op), + name="segment_mean.generic", + ) + return strategy + + def wrap_compute_unique(topi_compute): """Wrap unique topi compute""" @@ -1594,8 +1663,54 @@ def unique_strategy(attrs, inputs, out_type, target): return strategy +# segment_sum +def wrap_compute_segment_sum(topi_compute): + """wrap segment_sum topi compute""" + + def _compute_segment_sum(attrs, inputs, out_type): + num_segments = attrs.num_segments + return [topi_compute(inputs[0], inputs[1], num_segments, "sum")] + + return _compute_segment_sum + + +@override_native_generic_func("segment_sum_strategy") +def segment_sum_strategy(attrs, inputs, out_type, target): + """segment_sum generic strategy""" + strategy = _op.OpStrategy() + strategy.add_implementation( + wrap_compute_segment_sum(topi.segment_op), + wrap_topi_schedule(topi.generic.schedule_segment_op), + name="segment_sum.generic", + ) + return strategy + + @generic_func def schedule_transpose(attrs, outs, target): """schedule transpose""" with target: return schedule_injective(attrs, outs, target) + + +# segment_prod +def wrap_compute_segment_prod(topi_compute): + """wrap segment_prod topi compute""" + + def _compute_segment_prod(attrs, inputs, out_type): + num_segments = attrs.num_segments + return [topi_compute(inputs[0], inputs[1], num_segments, "prod")] + + return _compute_segment_prod + + +@override_native_generic_func("segment_prod_strategy") +def segment_prod_strategy(attrs, inputs, out_type, target): + """segment_prod generic strategy""" + strategy = _op.OpStrategy() + strategy.add_implementation( + wrap_compute_segment_prod(topi.segment_op), + wrap_topi_schedule(topi.generic.schedule_segment_op), + name="segment_prod.generic", + ) + return strategy diff --git a/python/tvm/relay/op/tensor.py b/python/tvm/relay/op/tensor.py index a38a23064d6f..259bd2a0d032 100644 --- a/python/tvm/relay/op/tensor.py +++ b/python/tvm/relay/op/tensor.py @@ -1285,3 +1285,113 @@ def isinf(data): The computed result. """ return _make.isinf(data) + + +def segment_max(data, segment_ids, num_segments): + """Computes the maximum along segments of a tensor. + + Parameters + ---------- + data : relay.Expr + The input data + + segment_ids : relay.Expr + The segments data + + num_segments : int + The maximum of segment_ids. + + Returns + ------- + result : relay.Expr + The computed result. + """ + return _make.segment_max(data, segment_ids, num_segments) + + +def segment_min(data, segment_ids, num_segments): + """Computes the minimum along segments of a tensor. + + Parameters + ---------- + data : relay.Expr + The input data + + segment_ids : relay.Expr + The segments data + + num_segments : int + The maximum of segment_ids. + + Returns + ------- + result : relay.Expr + The computed result. + """ + return _make.segment_min(data, segment_ids, num_segments) + + +def segment_mean(data, segment_ids, num_segments): + """Computes the mean along segments of a tensor. + + Parameters + ---------- + data : relay.Expr + The input data + + segment_ids : relay.Expr + The segments data + + num_segments : int + The maximum of segment_ids. + + Returns + ------- + result : relay.Expr + The computed result. + """ + return _make.segment_mean(data, segment_ids, num_segments) + + +def segment_sum(data, segment_ids, num_segments): + """Computes the sum along segments of a tensor. + + Parameters + ---------- + data : relay.Expr + The input data + + segment_ids : relay.Expr + The segments data + + num_segments : int + The maximum of segment_ids. + + Returns + ------- + result : relay.Expr + The computed result. + """ + return _make.segment_sum(data, segment_ids, num_segments) + + +def segment_prod(data, segment_ids, num_segments): + """Computes the prod along segments of a tensor. + + Parameters + ---------- + data : relay.Expr + The input data + + segment_ids : relay.Expr + The segments data + + num_segments : int + The maximum of segment_ids. + + Returns + ------- + result : relay.Expr + The computed result. + """ + return _make.segment_prod(data, segment_ids, num_segments) diff --git a/python/tvm/topi/generic/vision.py b/python/tvm/topi/generic/vision.py index e7518b1110a1..751fb78a31a4 100644 --- a/python/tvm/topi/generic/vision.py +++ b/python/tvm/topi/generic/vision.py @@ -176,3 +176,20 @@ def schedule_proposal(outs): The computation schedule for the op. """ return _default_schedule(outs, False) + + +def schedule_segment_op(outs): + """Schedule for segment_max operator. + + Parameters + ---------- + outs: Array of Tensor + The computation graph description of segment_max + in the format of an array of tensors. + + Returns + ------- + s: Schedule + The computation schedule for the op. + """ + return _default_schedule(outs, False) diff --git a/python/tvm/topi/tensor.py b/python/tvm/topi/tensor.py index 31ebe86760cb..1886554e064f 100644 --- a/python/tvm/topi/tensor.py +++ b/python/tvm/topi/tensor.py @@ -18,6 +18,8 @@ """Elementwise operators""" from __future__ import absolute_import as _abs from . import cpp +from .. import te, tir +from .utils import get_const_tuple def elemwise_sum(xs): @@ -73,3 +75,104 @@ def full_like(x, fill_value): The result. """ return cpp.full_like(x, fill_value) + + +def segment_op(data, segment_ids, num_out, opname): + """segmet_max, segmet_min, segmet_mean, segmet_sum, segmet_prod + unsorted_segment_max, unsorted_segmet_min, unsorted_segmet_mean, + unsorted_segmet_sum, unsorted_segmet_prod operators. + + Parameters + ---------- + data : tvm.te.Tensor + input data + + segment_ids : tvm.te.Tensor + input segment ids + + num_out : int + number of output + + opname : str + name of target op + + Returns + ------- + out : tvm.te.Tensor + Tensor with shape determined by the segment ids. + """ + + def _max(in_data, out): + return te.max(in_data, out) + + def _min(in_data, out): + return te.min(in_data, out) + + def _add(in_data, out): + return in_data + out + + def _prod(in_data, out): + return in_data * out + + func_dict = {"max": _max, "min": _min, "mean": _add, "sum": _add, "prod": _prod} + init_dict = { + "max": -float("inf"), + "min": float("inf"), + "mean": 0.0, + "sum": 0.0, + "prod": 1.0, + } + + def _segment_op(data, segment_ids, out_buf): + func = func_dict[opname] + init_data = init_dict[opname] + + ib = tir.ir_builder.create() + input_data = ib.buffer_ptr(data) + seg_ids = ib.buffer_ptr(segment_ids) + out = ib.buffer_ptr(out_buf) + num = ib.allocate("int32", (1), name="num", scope="local") + + shape = get_const_tuple(data.shape) + num_segment = get_const_tuple(out_buf.shape)[0] + # The number of data to be calculated for each output + inner_size = 1 + for s in range(1, len(shape)): + inner_size = inner_size * shape[s] + + # Init output tensor + with ib.for_range(0, num_segment) as n: + with ib.for_range(0, inner_size) as j: + out_index = n * inner_size + j + out[out_index] = init_data + + num[0] = 0 + # Operate on numbers with the same id + with ib.for_range(0, shape[0]) as k: + with ib.if_scope(seg_ids[k] == n): + with ib.for_range(0, inner_size) as l: + out_index = n * inner_size + l + in_index = k * inner_size + l + out[out_index] = func(input_data[in_index], out[out_index]) + num[0] += 1 + if opname == "mean": + with ib.if_scope(num[0] > 0): + with ib.for_range(0, inner_size) as s: + out_index = n * inner_size + s + out[out_index] = out[out_index] / num[0] + + return ib.get() + + assert len(segment_ids.shape) == 1 + + out_shape = list(get_const_tuple(data.shape)) + out_shape[0] = num_out + + out = te.extern( + out_shape, + [data, segment_ids], + lambda ins, outs: _segment_op(ins[0], ins[1], outs[0]), + dtype=data.dtype, + ) + + return out diff --git a/src/relay/op/tensor/binary.cc b/src/relay/op/tensor/binary.cc index aafd4492fec4..1513b063c9ea 100644 --- a/src/relay/op/tensor/binary.cc +++ b/src/relay/op/tensor/binary.cc @@ -161,5 +161,115 @@ RELAY_REGISTER_CMP_OP("greater_equal") .set_support_level(4) .set_attr("FTVMCompute", RELAY_BINARY_COMPUTE(topi::greater_equal)); +// segment_max +TVM_REGISTER_NODE_TYPE(SegmentAttrs); + +bool SegmentRel(const Array& types, int num_inputs, const Attrs& attrs, + const TypeReporter& reporter) { + CHECK_EQ(types.size(), 3); + + auto segment_attrs = attrs.as(); + int num_segments = segment_attrs->num_segments; + + const auto* data = types[0].as(); + const auto& dshape = data->shape; + Array oshape(dshape); + oshape.Set(0, num_segments); + + // assign output type + reporter->Assign(types[2], TensorType(oshape, data->dtype)); + return true; +} + +Expr MakeSegmentMax(Expr data, Expr segment_ids, int num_segments) { + auto attrs = make_object(); + attrs->num_segments = num_segments; + static const Op& op = Op::Get("segment_max"); + return Call(op, {data, segment_ids}, Attrs(attrs), {}); +} + +TVM_REGISTER_GLOBAL("relay.op._make.segment_max").set_body_typed(MakeSegmentMax); + +RELAY_REGISTER_OP("segment_max") + .describe(R"doc(Computes the maximum along segments of a tensor. +)doc" TVM_ADD_FILELINE) + .set_num_inputs(2) + .add_argument("data", "Tensor", "Input data.") + .add_argument("segment_ids", "Tensor", "Segments tensor.") + .set_support_level(5) + .add_type_rel("SegmentMax", SegmentRel); + +Expr MakeSegmentMin(Expr data, Expr segment_ids, int num_segments) { + auto attrs = make_object(); + attrs->num_segments = num_segments; + static const Op& op = Op::Get("segment_min"); + return Call(op, {data, segment_ids}, Attrs(attrs), {}); +} + +TVM_REGISTER_GLOBAL("relay.op._make.segment_min").set_body_typed(MakeSegmentMin); + +RELAY_REGISTER_OP("segment_min") + .describe(R"doc(Computes the minimum along segments of a tensor. +)doc" TVM_ADD_FILELINE) + .set_num_inputs(2) + .add_argument("data", "Tensor", "Input data.") + .add_argument("segment_ids", "Tensor", "Segments tensor.") + .set_support_level(5) + .add_type_rel("SegmentMin", SegmentRel); + +Expr MakeSegmentMean(Expr data, Expr segment_ids, int num_segments) { + auto attrs = make_object(); + attrs->num_segments = num_segments; + static const Op& op = Op::Get("segment_mean"); + return Call(op, {data, segment_ids}, Attrs(attrs), {}); +} + +TVM_REGISTER_GLOBAL("relay.op._make.segment_mean").set_body_typed(MakeSegmentMean); + +RELAY_REGISTER_OP("segment_mean") + .describe(R"doc(Computes the mean along segments of a tensor. +)doc" TVM_ADD_FILELINE) + .set_num_inputs(2) + .add_argument("data", "Tensor", "Input data.") + .add_argument("segment_ids", "Tensor", "Segments tensor.") + .set_support_level(5) + .add_type_rel("SegmentMean", SegmentRel); + +Expr MakeSegmentSum(Expr data, Expr segment_ids, int num_segments) { + auto attrs = make_object(); + attrs->num_segments = num_segments; + static const Op& op = Op::Get("segment_sum"); + return Call(op, {data, segment_ids}, Attrs(attrs), {}); +} + +TVM_REGISTER_GLOBAL("relay.op._make.segment_sum").set_body_typed(MakeSegmentSum); + +RELAY_REGISTER_OP("segment_sum") + .describe(R"doc(Computes the sum along segments of a tensor. +)doc" TVM_ADD_FILELINE) + .set_num_inputs(2) + .add_argument("data", "Tensor", "Input data.") + .add_argument("segment_ids", "Tensor", "Segments tensor.") + .set_support_level(5) + .add_type_rel("SegmentSum", SegmentRel); + +Expr MakeSegmentProd(Expr data, Expr segment_ids, int num_segments) { + auto attrs = make_object(); + attrs->num_segments = num_segments; + static const Op& op = Op::Get("segment_prod"); + return Call(op, {data, segment_ids}, Attrs(attrs), {}); +} + +TVM_REGISTER_GLOBAL("relay.op._make.segment_prod").set_body_typed(MakeSegmentProd); + +RELAY_REGISTER_OP("segment_prod") + .describe(R"doc(Computes the prod along segments of a tensor. +)doc" TVM_ADD_FILELINE) + .set_num_inputs(2) + .add_argument("data", "Tensor", "Input data.") + .add_argument("segment_ids", "Tensor", "Segments tensor.") + .set_support_level(5) + .add_type_rel("SegmentProd", SegmentRel); + } // namespace relay } // namespace tvm diff --git a/tests/python/frontend/tensorflow/test_forward.py b/tests/python/frontend/tensorflow/test_forward.py index f29450dbb604..203d5ba479ee 100644 --- a/tests/python/frontend/tensorflow/test_forward.py +++ b/tests/python/frontend/tensorflow/test_forward.py @@ -4398,6 +4398,95 @@ def _test_forward_rint(shape): _test_forward_rint([2, 5, 2, 5]) +def test_forward_segment(): + """test operator segmet_max, segmet_min, segmet_mean, segmet_sum, segmet_prod """ + + def _get_segment_ids(length, size): + segment_ids = [0] + for i in range(size): + if np.array(segment_ids).sum() < length: + for _ in range(i): + segment_ids.append(i) + else: + break + length = length - len(segment_ids) + for i in range(length): + segment_ids.append(segment_ids[-1]) + return np.array(segment_ids).astype("int32") + + def _test_forward_segment(name, data_shape, size): + np_data = np.random.uniform(-100, 100, size=data_shape).astype(np.float32) + tf.reset_default_graph() + in_data = tf.placeholder(tf.float32, data_shape, name="in_data") + segment_ids = _get_segment_ids(data_shape[0], size) + tf_segment_ids = tf.constant(segment_ids) + segment_op = getattr(tf.math, name) + result = segment_op(in_data, tf_segment_ids, name=name) + compare_tf_with_tvm([np_data, segment_ids], ["in_data:0"], result.name, no_gpu=True) + + _test_forward_segment("segment_max", [100], 20) + _test_forward_segment("segment_max", [10, 3, 4], 5) + _test_forward_segment("segment_max", [10, 3, 4, 4], 10) + _test_forward_segment("segment_min", [100], 20) + _test_forward_segment("segment_min", [10, 3, 4], 5) + _test_forward_segment("segment_min", [10, 3, 4, 4], 10) + _test_forward_segment("segment_mean", [100], 20) + _test_forward_segment("segment_mean", [10, 3, 4], 5) + _test_forward_segment("segment_mean", [10, 3, 4, 4], 10) + _test_forward_segment("segment_sum", [100], 20) + _test_forward_segment("segment_sum", [10, 3, 4], 5) + _test_forward_segment("segment_sum", [10, 3, 4, 4], 10) + _test_forward_segment("segment_prod", [100], 20) + _test_forward_segment("segment_prod", [10, 3, 4], 5) + _test_forward_segment("segment_prod", [10, 3, 4, 4], 10) + + +def test_forward_unsorted_segment(): + """test operator unsorted_segment_max, unsorted_segmet_min, + unsorted_segmet_mean, unsorted_segmet_sum, unsorted_segmet_prod""" + + def _get_segment_ids(length, size): + segment_ids = [0] + for i in range(size): + if np.array(segment_ids).sum() < length: + for _ in range(i): + segment_ids.append(i) + else: + break + length = length - len(segment_ids) + for i in range(length): + segment_ids.append(segment_ids[-1]) + np.random.shuffle(segment_ids) + segment_ids = np.array(segment_ids).astype("int32") + return segment_ids, segment_ids.max() + 1 + + def _test_forward_unsorted_segment(name, data_shape, size): + np_data = np.random.uniform(-100, 100, size=data_shape).astype(np.float32) + tf.reset_default_graph() + in_data = tf.placeholder(tf.float32, data_shape, name="in_data") + segment_ids, num_segments = _get_segment_ids(data_shape[0], size) + tf_segment_ids = tf.constant(segment_ids) + segment_op = getattr(tf.math, name) + result = segment_op(in_data, tf_segment_ids, num_segments, name=name) + compare_tf_with_tvm([np_data, segment_ids], ["in_data:0"], result.name, no_gpu=True) + + _test_forward_unsorted_segment("unsorted_segment_max", [100], 20) + _test_forward_unsorted_segment("unsorted_segment_max", [10, 3, 4], 5) + _test_forward_unsorted_segment("unsorted_segment_max", [10, 3, 4, 4], 10) + _test_forward_unsorted_segment("unsorted_segment_min", [100], 20) + _test_forward_unsorted_segment("unsorted_segment_min", [10, 3, 4], 5) + _test_forward_unsorted_segment("unsorted_segment_min", [10, 3, 4, 4], 10) + _test_forward_unsorted_segment("unsorted_segment_mean", [100], 20) + _test_forward_unsorted_segment("unsorted_segment_mean", [10, 3, 4], 5) + _test_forward_unsorted_segment("unsorted_segment_mean", [10, 3, 4, 4], 10) + _test_forward_unsorted_segment("unsorted_segment_sum", [100], 20) + _test_forward_unsorted_segment("unsorted_segment_sum", [10, 3, 4], 5) + _test_forward_unsorted_segment("unsorted_segment_sum", [10, 3, 4, 4], 10) + _test_forward_unsorted_segment("unsorted_segment_prod", [100], 20) + _test_forward_unsorted_segment("unsorted_segment_prod", [10, 3, 4], 5) + _test_forward_unsorted_segment("unsorted_segment_prod", [10, 3, 4, 4], 10) + + def test_forward_negative(): """test tf operator Neg""" np_data = np.random.uniform(-100, 255, size=(224, 224, 3)).astype(np.float32) diff --git a/tests/python/topi/python/test_topi_math.py b/tests/python/topi/python/test_topi_math.py index c7f80033bdf3..5ba95b3c5295 100644 --- a/tests/python/topi/python/test_topi_math.py +++ b/tests/python/topi/python/test_topi_math.py @@ -234,8 +234,156 @@ def check_target(target): ) +def verify_segment(name, np_data, np_segment_ids, np_out): + device = "llvm" + ctx = tvm.context(device, 0) + print("Running on target: %s" % device) + data = te.placeholder(np_data.shape) + segment_ids = te.placeholder(np_segment_ids.shape, dtype="int32") + num_out = np_segment_ids[-1] + 1 + with tvm.target.Target(device): + fcompute, fschedule = tvm.topi.testing.dispatch( + device, + { + "generic": (topi.segment_op, topi.generic.schedule_segment_op), + }, + ) + out = fcompute(data, segment_ids, num_out, name) + s = fschedule(out) + f = tvm.build(s, [data, segment_ids, out], device) + tvm_data = tvm.nd.array(np_data, ctx=ctx) + tvm_segment_ids = tvm.nd.array(np_segment_ids, ctx=ctx) + tvm_out = tvm.nd.empty(ctx=ctx, shape=out.shape, dtype=out.dtype) + f(tvm_data, tvm_segment_ids, tvm_out) + tvm.testing.assert_allclose(tvm_out.asnumpy(), np_out, rtol=1e-4) + + +def test_segment(): + # segmet_max, segmet_min, segmet_mean, segmet_sum, segmet_prod + np_data = np.array([0, 0.8, 1, 20, -25, 45, 1, 0.7, -30, 60, 50, 80]).astype("float32") + segment_ids = np.array([0, 1, 1, 2, 2, 2, 3, 4, 5, 5, 5, 6]).astype("int32") + np_result = np.array([0, 1, 45, 1, 0.7, 60, 80]) + verify_segment("max", np_data, segment_ids, np_result) + + np_data = np.array( + [ + [0, 0.8, 1, 20, -25, 45], + [1, 0.7, 30, 60, 50, 80], + [0, 0.4, 4, 21, 19, 40], + [2, -0.9, 35, 61, 52, 79], + [1, 0.5, 100, 60, 70, 110], + ] + ).astype("float32") + segment_ids = np.array([0, 0, 1, 1, 2]).astype("int32") + np_result = np.array( + [ + [1, 0.8, 30, 60, 50, 80], + [2, 0.4, 35, 61, 52, 79], + [1, 0.5, 100, 60, 70, 110], + ] + ) + verify_segment("max", np_data, segment_ids, np_result) + + np_data = np.array([0, 0.8, 1, 20, -25, 45, 1, 0.7, -30, 60, 50, 80]).astype("float32") + segment_ids = np.array([0, 1, 1, 2, 2, 2, 3, 4, 5, 5, 5, 6]).astype("int32") + np_result = np.array([0, 0.8, -25, 1, 0.7, -30, 80]) + verify_segment("min", np_data, segment_ids, np_result) + + np_data = np.array( + [ + [0, 0.8, 1, 20, -25, 45], + [1, 0.7, 30, 60, 50, 80], + [0, 0.4, 4, 21, 19, 40], + [2, -0.9, 35, 61, 52, 79], + [1, 0.5, 100, 60, 70, 110], + ] + ).astype("float32") + segment_ids = np.array([0, 0, 1, 1, 2]).astype("int32") + np_result = np.array( + [ + [0.0, 0.7, 1.0, 20.0, -25.0, 45.0], + [0.0, -0.9, 4.0, 21.0, 19.0, 40.0], + [1.0, 0.5, 100.0, 60.0, 70.0, 110.0], + ] + ) + verify_segment("min", np_data, segment_ids, np_result) + + np_data = np.array([0, 0.8, 1, 20, -25, 45, 1, 0.7, -30, 60, 50, 80]).astype("float32") + segment_ids = np.array([0, 1, 1, 2, 2, 2, 3, 4, 5, 5, 5, 6]).astype("int32") + np_result = np.array([0.0, 0.9, 13.333333, 1.0, 0.7, 26.666666, 80.0]) + verify_segment("mean", np_data, segment_ids, np_result) + + np_data = np.array( + [ + [0, 0.8, 1, 20, -25, 45], + [1, 0.7, 30, 60, 50, 80], + [0, 0.4, 4, 21, 19, 40], + [2, -0.9, 35, 61, 52, 79], + [1, 0.5, 100, 60, 70, 110], + ] + ).astype("float32") + segment_ids = np.array([0, 0, 1, 1, 2]).astype("int32") + np_result = np.array( + [ + [0.5, 0.75, 15.5, 40.0, 12.5, 62.5], + [1.0, -0.25, 19.5, 41.0, 35.5, 59.5], + [1.0, 0.5, 100.0, 60.0, 70.0, 110.0], + ] + ) + verify_segment("mean", np_data, segment_ids, np_result) + + np_data = np.array([0, 0.8, 1, 20, -25, 45, 1, 0.7, -30, 60, 50, 80]).astype("float32") + segment_ids = np.array([0, 1, 1, 2, 2, 2, 3, 4, 5, 5, 5, 6]).astype("int32") + np_result = np.array([0.0, 1.8, 40.0, 1.0, 0.7, 80.0, 80.0]) + verify_segment("sum", np_data, segment_ids, np_result) + + np_data = np.array( + [ + [0, 0.8, 1, 20, -25, 45], + [1, 0.7, 30, 60, 50, 80], + [0, 0.4, 4, 21, 19, 40], + [2, -0.9, 35, 61, 52, 79], + [1, 0.5, 100, 60, 70, 110], + ] + ).astype("float32") + segment_ids = np.array([0, 0, 1, 1, 2]).astype("int32") + np_result = np.array( + [ + [1.0, 1.5, 31.0, 80.0, 25.0, 125.0], + [2.0, -0.5, 39.0, 82.0, 71.0, 119.0], + [1.0, 0.5, 100.0, 60.0, 70.0, 110.0], + ] + ) + verify_segment("sum", np_data, segment_ids, np_result) + + np_data = np.array([0, 0.8, 1, 20, -25, 45, 1, 0.7, -30, 60, 50, 80]).astype("float32") + segment_ids = np.array([0, 1, 1, 2, 2, 2, 3, 4, 5, 5, 5, 6]).astype("int32") + np_result = np.array([0.0, 0.8, -22500.0, 1.0, 0.7, -90000, 80]) + verify_segment("prod", np_data, segment_ids, np_result) + + np_data = np.array( + [ + [0, 0.8, 1, 20, -25, 45], + [1, 0.7, 30, 60, 50, 80], + [0, 0.4, 4, 21, 19, 40], + [2, -0.9, 35, 61, 52, 79], + [1, 0.5, 100, 60, 70, 110], + ] + ).astype("float32") + segment_ids = np.array([0, 0, 1, 1, 2]).astype("int32") + np_result = np.array( + [ + [0.0, 0.56, 30.0, 1200.0, -1250.0, 3600.0], + [0.0, -0.36, 140.0, 1281.0, 988.0, 3160.0], + [1.0, 0.5, 100.0, 60.0, 70, 110], + ] + ) + verify_segment("prod", np_data, segment_ids, np_result) + + if __name__ == "__main__": test_util() test_ewise() test_cast() test_fastmath() + test_segment()