From 3cfe1279e3d3a0297e7876cb02116f38f4c8e1f7 Mon Sep 17 00:00:00 2001 From: CircleSpin Date: Fri, 10 Sep 2021 17:35:27 -0400 Subject: [PATCH 01/11] adding dynamic squeeze first steps --- python/tvm/relay/op/transform.py | 6 +- src/relay/op/dyn/tensor/transform.cc | 117 ++++++++++++++++++ .../relay/dyn/test_dynamic_op_level3.py | 20 +++ 3 files changed, 142 insertions(+), 1 deletion(-) diff --git a/python/tvm/relay/op/transform.py b/python/tvm/relay/op/transform.py index fe1a73ca231a..234e76b11813 100644 --- a/python/tvm/relay/op/transform.py +++ b/python/tvm/relay/op/transform.py @@ -149,7 +149,7 @@ def squeeze(data, axis=None): data : tvm.relay.Expr The input data to the operator. - axis : None or List[int] + axis : None or List[int] or Expr The set of axes to remove. If axis = None, remove all axis of dimensions 1. If any specified axis has dimension that does not equal 1, it is an error. @@ -159,6 +159,10 @@ def squeeze(data, axis=None): result : tvm.relay.Expr The squeezed result. """ + if isinstance(axis, Constant): + axis = list(axis.data.numpy()) + if isinstance(axis, Expr): + return _dyn_make.squeeze(data, axis) return _make.squeeze(data, axis) diff --git a/src/relay/op/dyn/tensor/transform.cc b/src/relay/op/dyn/tensor/transform.cc index 848d058f0af3..c200457104ad 100644 --- a/src/relay/op/dyn/tensor/transform.cc +++ b/src/relay/op/dyn/tensor/transform.cc @@ -692,6 +692,123 @@ RELAY_REGISTER_OP("dyn.expand_dims") .set_attr("FTVMCompute", ExpandDimsCompute) .set_attr("TOpPattern", kInjective); +Expr MakeDynSqueeze(Expr data, Expr axes) { + auto attrs = make_object(); + static const Op& op = Op::Get("dyn.squeeze"); + return Call(op, {data, axes}, Attrs(attrs), {}); +} + +TVM_REGISTER_GLOBAL("relay.op.dyn._make.squeeze").set_body_typed(MakeDynSqueeze); + +bool DynSqueezeRel(const Array& types, int num_inputs, const Attrs& attrs, + const TypeReporter& reporter) { + ICHECK_EQ(types.size(), 3); + const auto* data = types[0].as(); + if (data == nullptr) { + return false; + } + const auto* axes = types[1].as(); + if (axes == nullptr) { + return false; + } + ICHECK_EQ(axes->shape.size(), 1); + ICHECK(axes->shape[0].as()); + size_t output_rank = data->shape.size() - axes->shape[0].as()->value; + std::vector result_shape(output_rank, Any()); + reporter->Assign(types[2], TensorType(result_shape, data->dtype)); + return true; +} + + +// Array SqueezeCompute(const Attrs& attrs, const Array& inputs, +// const Type& out_type) { +// const SqueezeAttrs* param = attrs.as(); +// ICHECK(param != nullptr); +// return {topi::squeeze(inputs[0], param->axis)}; +// } + +// InferCorrectLayoutOutput SqueezeInferCorrectLayout(const Attrs& attrs, +// const Array& new_in_layouts, +// const Array& old_in_layouts, +// const Array& old_in_types) { +// const auto* attrs_ptr = attrs.as(); +// ICHECK(attrs_ptr); +// ObjectPtr params = make_object(*attrs_ptr); + +// Layout inferred_input = new_in_layouts.defined() ? new_in_layouts[0] : old_in_layouts[0]; +// Layout inferred_output = inferred_input; + +// ICHECK(old_in_types[0].as()); +// const auto& shape = old_in_types[0].as()->shape; + +// // axis to squeeze +// Array axis; +// if (params->axis.defined()) { +// axis = params->axis; +// } else { +// // if axes is None, squeeze all axes of dimension 1 +// for (size_t i = 0; i < shape.size(); i++) { +// if (topi::detail::GetConstInt(shape[i]) == 1) { +// axis.push_back(i); +// } +// } +// } + +// // If new_in_layouts are defined, this code tries to modify the layout +// if (new_in_layouts.defined() && old_in_layouts.defined()) { +// Array new_axis; +// for (const auto& e : axis) { +// const auto& dim = old_in_layouts[0][e]; +// new_axis.push_back((new_in_layouts[0]).IndexOf(dim)); +// } +// params->axis = new_axis; +// axis = new_axis; +// } + +// // Infer output layout +// Array kept_axes; +// for (size_t i = 0; i < inferred_input.ndim(); i++) { +// bool is_dim_kept = true; + +// // Check whether the dim should be kept +// for (const auto& e : axis) { +// int64_t axis_val = e->value; +// if (axis_val < 0) { +// axis_val += inferred_input.ndim(); +// } +// if (static_cast(i) == axis_val) { +// is_dim_kept = false; +// break; +// } +// } + +// if (is_dim_kept) { +// kept_axes.push_back(inferred_input->axes[i]); +// } +// } +// inferred_output = Layout(kept_axes); + +// return InferCorrectLayoutOutput({inferred_input}, {inferred_output}, Attrs(params)); +// } + +RELAY_REGISTER_OP("dyn.squeeze") + .describe(R"code(Squeeze the input tensor at the dimensions given by axes + +- **data**: The input data to the operator. +- **axes**: The axes to squeeze. + +)code" TVM_ADD_FILELINE) + .set_num_inputs(2) + .set_attrs_type() + .add_argument("data", "Tensor", "The input tensor.") + .add_argument("axes", "Tensor", "The axes to squeeze.") + .set_support_level(3) + .add_type_rel("DynSqueeze", DynSqueezeRel); + // .set_attr("FTVMCompute", SqueezeCompute) + // .set_attr("TOpPattern", kInjective) + // .set_attr("FInferCorrectLayout", SqueezeInferCorrectLayout) + // .set_attr("TReshapeOp", true); + } // namespace dyn } // namespace relay } // namespace tvm diff --git a/tests/python/relay/dyn/test_dynamic_op_level3.py b/tests/python/relay/dyn/test_dynamic_op_level3.py index 8c57e1dc4a9f..4aa8e6bfc3dc 100644 --- a/tests/python/relay/dyn/test_dynamic_op_level3.py +++ b/tests/python/relay/dyn/test_dynamic_op_level3.py @@ -92,6 +92,26 @@ def verify_reshape(shape, newshape, oshape): verify_reshape((4, 7), (2, 7, 2), (2, 7, 2)) +def test_squeeze(): + def verify_squeeze(shape, dtype, axis): + x = relay.var("x", relay.TensorType(shape, dtype)) + if axis is not None: + axis = relay.var("axis", relay.TensorType([len(axis)], "int64")) + squeeze = relay.squeeze(x, axis=axis) + print(squeeze) + print(run_infer_type(squeeze)) + + # np_axis = tuple(axis) if axis is not None else None + + # data = np.random.random_sample(shape).astype(dtype) + # op_res = create_executor().evaluate(squeeze, {x: relay.const(data)}) + # ref_res = np.squeeze(data, axis=np_axis) + # np.testing.assert_allclose(op_res.numpy(), ref_res, rtol=0.01) + + verify_squeeze((1, 3, 2, 5), "float32", None) + verify_squeeze((1, 3, 1), "float32", [0]) + verify_squeeze((1, 2, 1, 2, 1), "float32", [0, 2]) + @tvm.testing.uses_gpu def test_dyn_expand_dims(): def verify_expand_dims( From dcba2fe089de9337e2811acbdd11848149452c8a Mon Sep 17 00:00:00 2001 From: CircleSpin Date: Mon, 20 Sep 2021 17:27:13 -0400 Subject: [PATCH 02/11] Matt B. implementing shape --- python/tvm/relay/op/_transform.py | 29 +++++++++++ src/relay/op/dyn/tensor/transform.cc | 52 ++++++++++++------- .../relay/dyn/test_dynamic_op_level3.py | 21 ++++---- 3 files changed, 73 insertions(+), 29 deletions(-) diff --git a/python/tvm/relay/op/_transform.py b/python/tvm/relay/op/_transform.py index 0284d2483ce5..fdde8230d53c 100644 --- a/python/tvm/relay/op/_transform.py +++ b/python/tvm/relay/op/_transform.py @@ -1174,3 +1174,32 @@ def gather_nd_shape_func(attrs, inputs, _): assert index_rank > 0, "index_rank needs to be specified for dynamic gather_nd" return [_gather_nd_shape(inputs[0], inputs[1], convert(batch_dims), convert(index_rank))] + + +@script +def _squeeze_shape_func_input_data(data, axis, ndims): + out = output_tensor((ndims,), "int64") + out_i = 0 + for i in const_range(data.shape[0]): + not_in_axis = True + for j in const_range(axis.shape[0]): + if i == axis[j]: + not_in_axis = False + if not_in_axis: + out[out_i] = int64(data[i]) + out_i += 1 + + return out + + +# def _squeeze_shape_func_input_data(data, axis, ndims): +# out = [] +# for i in range(len(data)): +# if not i in axis: +# out.append(data[i]) +# return out + + +@_reg.register_shape_func("dyn.squeeze", [False, True]) +def dynamic_squeeze_shape_func(attrs, inputs, out_ndims): + return [_squeeze_shape_func_input_data(inputs[0], inputs[1], out_ndims[0])] \ No newline at end of file diff --git a/src/relay/op/dyn/tensor/transform.cc b/src/relay/op/dyn/tensor/transform.cc index c200457104ad..d099432909c7 100644 --- a/src/relay/op/dyn/tensor/transform.cc +++ b/src/relay/op/dyn/tensor/transform.cc @@ -692,14 +692,6 @@ RELAY_REGISTER_OP("dyn.expand_dims") .set_attr("FTVMCompute", ExpandDimsCompute) .set_attr("TOpPattern", kInjective); -Expr MakeDynSqueeze(Expr data, Expr axes) { - auto attrs = make_object(); - static const Op& op = Op::Get("dyn.squeeze"); - return Call(op, {data, axes}, Attrs(attrs), {}); -} - -TVM_REGISTER_GLOBAL("relay.op.dyn._make.squeeze").set_body_typed(MakeDynSqueeze); - bool DynSqueezeRel(const Array& types, int num_inputs, const Attrs& attrs, const TypeReporter& reporter) { ICHECK_EQ(types.size(), 3); @@ -720,12 +712,36 @@ bool DynSqueezeRel(const Array& types, int num_inputs, const Attrs& attrs, } -// Array SqueezeCompute(const Attrs& attrs, const Array& inputs, -// const Type& out_type) { -// const SqueezeAttrs* param = attrs.as(); -// ICHECK(param != nullptr); -// return {topi::squeeze(inputs[0], param->axis)}; -// } +Array SqueezeCompute(const Attrs& attrs, const Array& inputs, + const Type& out_type) { + const auto* out_ttype = out_type.as(); + ICHECK(out_ttype != nullptr); + Array newshape; + for (auto val : out_ttype->shape) { + if (val->IsInstance()) { + newshape.push_back(val.as()->ToVar()); + } else { + newshape.push_back(val); + } + } + // newshape = {1, shape_var0} + return {topi::reshape(inputs[0], newshape)}; + // def dyn_reshape(x, shape_var0): + // out = zeros((1 * shape_var0,)) + // for j in range(1): + // for k in range(shape_var0): + // out[j*shape_var0 + k] = x[j * shape_var0 + k] + // return out +} + +Expr MakeDynSqueeze(Expr data, Expr axes) { + auto attrs = make_object(); + static const Op& op = Op::Get("dyn.squeeze"); + return Call(op, {data, axes}, Attrs(attrs), {}); +} + +TVM_REGISTER_GLOBAL("relay.op.dyn._make.squeeze").set_body_typed(MakeDynSqueeze); + // InferCorrectLayoutOutput SqueezeInferCorrectLayout(const Attrs& attrs, // const Array& new_in_layouts, @@ -803,11 +819,11 @@ RELAY_REGISTER_OP("dyn.squeeze") .add_argument("data", "Tensor", "The input tensor.") .add_argument("axes", "Tensor", "The axes to squeeze.") .set_support_level(3) - .add_type_rel("DynSqueeze", DynSqueezeRel); - // .set_attr("FTVMCompute", SqueezeCompute) - // .set_attr("TOpPattern", kInjective) + .add_type_rel("DynSqueeze", DynSqueezeRel) + .set_attr("FTVMCompute", SqueezeCompute) + .set_attr("TOpPattern", kInjective) // .set_attr("FInferCorrectLayout", SqueezeInferCorrectLayout) - // .set_attr("TReshapeOp", true); + .set_attr("TReshapeOp", true); } // namespace dyn } // namespace relay diff --git a/tests/python/relay/dyn/test_dynamic_op_level3.py b/tests/python/relay/dyn/test_dynamic_op_level3.py index 4aa8e6bfc3dc..f33743b9be02 100644 --- a/tests/python/relay/dyn/test_dynamic_op_level3.py +++ b/tests/python/relay/dyn/test_dynamic_op_level3.py @@ -95,20 +95,19 @@ def verify_reshape(shape, newshape, oshape): def test_squeeze(): def verify_squeeze(shape, dtype, axis): x = relay.var("x", relay.TensorType(shape, dtype)) - if axis is not None: - axis = relay.var("axis", relay.TensorType([len(axis)], "int64")) + assert axis is not None + np_axis = tuple(axis) + axis = relay.var("axis", relay.TensorType([len(axis)], "int64")) squeeze = relay.squeeze(x, axis=axis) - print(squeeze) - print(run_infer_type(squeeze)) - - # np_axis = tuple(axis) if axis is not None else None - # data = np.random.random_sample(shape).astype(dtype) - # op_res = create_executor().evaluate(squeeze, {x: relay.const(data)}) - # ref_res = np.squeeze(data, axis=np_axis) - # np.testing.assert_allclose(op_res.numpy(), ref_res, rtol=0.01) + print(squeeze) + func = relay.Function([x, axis], squeeze) + print(func) + + x_data = np.random.random_sample(shape).astype(dtype) + ref_res = np.squeeze(x_data, axis=np_axis) + verify_func(func, [x_data, np.array(np_axis).astype("int64")], ref_res) - verify_squeeze((1, 3, 2, 5), "float32", None) verify_squeeze((1, 3, 1), "float32", [0]) verify_squeeze((1, 2, 1, 2, 1), "float32", [0, 2]) From 9377ece8916fe7e1267f7d2039c12d3e54be1cf4 Mon Sep 17 00:00:00 2001 From: CircleSpin Date: Wed, 22 Sep 2021 10:26:04 -0400 Subject: [PATCH 03/11] squeeze implemented, dynamic_to_static and onnx importer next --- python/tvm/relay/frontend/onnx.py | 11 ++++ python/tvm/relay/op/_transform.py | 29 ---------- python/tvm/relay/op/dyn/_transform.py | 30 ++++++++++ src/relay/op/dyn/tensor/transform.cc | 66 ---------------------- tests/python/frontend/onnx/test_forward.py | 4 +- 5 files changed, 43 insertions(+), 97 deletions(-) diff --git a/python/tvm/relay/frontend/onnx.py b/python/tvm/relay/frontend/onnx.py index dfc0298979e6..49610d63ed8d 100644 --- a/python/tvm/relay/frontend/onnx.py +++ b/python/tvm/relay/frontend/onnx.py @@ -1482,7 +1482,18 @@ def _impl_v12(cls, inputs, attr, params): result = _op.expand_dims(result, axis) return result +class Squeeze(OnnxOpConverter): + """Operator converter for Squeeze.""" + @classmethod + def _impl_v1(cls, inputs, attr, params): + + @classmethod + def _impl_v13(cls, inputs, attr, params): + rank = _op.shape_of(_op.shape_of(inputs[0])) + axis = inputs[1] + axis = _op.where(axis < _op.const(0, *axis dtype*), axis + rank, axis) + return _op.squeeze(inputs[1], axis) class Split(OnnxOpConverter): """Operator converter for Split.""" diff --git a/python/tvm/relay/op/_transform.py b/python/tvm/relay/op/_transform.py index fdde8230d53c..0284d2483ce5 100644 --- a/python/tvm/relay/op/_transform.py +++ b/python/tvm/relay/op/_transform.py @@ -1174,32 +1174,3 @@ def gather_nd_shape_func(attrs, inputs, _): assert index_rank > 0, "index_rank needs to be specified for dynamic gather_nd" return [_gather_nd_shape(inputs[0], inputs[1], convert(batch_dims), convert(index_rank))] - - -@script -def _squeeze_shape_func_input_data(data, axis, ndims): - out = output_tensor((ndims,), "int64") - out_i = 0 - for i in const_range(data.shape[0]): - not_in_axis = True - for j in const_range(axis.shape[0]): - if i == axis[j]: - not_in_axis = False - if not_in_axis: - out[out_i] = int64(data[i]) - out_i += 1 - - return out - - -# def _squeeze_shape_func_input_data(data, axis, ndims): -# out = [] -# for i in range(len(data)): -# if not i in axis: -# out.append(data[i]) -# return out - - -@_reg.register_shape_func("dyn.squeeze", [False, True]) -def dynamic_squeeze_shape_func(attrs, inputs, out_ndims): - return [_squeeze_shape_func_input_data(inputs[0], inputs[1], out_ndims[0])] \ No newline at end of file diff --git a/python/tvm/relay/op/dyn/_transform.py b/python/tvm/relay/op/dyn/_transform.py index c8235ec9375a..1d855cb9420d 100644 --- a/python/tvm/relay/op/dyn/_transform.py +++ b/python/tvm/relay/op/dyn/_transform.py @@ -26,6 +26,7 @@ _reg.register_broadcast_schedule("dyn.broadcast_to") _reg.register_injective_schedule("dyn.reshape") _reg.register_injective_schedule("dyn.expand_dims") +_reg.register_injective_schedule("dyn.squeeze") _reg.register_broadcast_schedule("dyn.tile") _reg.register_injective_schedule("dyn.one_hot") _reg.register_injective_schedule("dyn.full") @@ -258,3 +259,32 @@ def _sparse_to_dense_shape_func(output_shape, ndim): @_reg.register_shape_func("dyn.sparse_to_dense", True) def sparse_to_dense_shape_func(attrs, inputs, out_ndims): return [_sparse_to_dense_shape_func(inputs[3], out_ndims[0])] + + +@script +def _squeeze_shape_func_input_data(data, axis, ndims): + out = output_tensor((ndims,), "int64") + out_i = 0 + for i in const_range(data.shape[0]): + not_in_axis = True + for j in const_range(axis.shape[0]): + if i == axis[j]: + not_in_axis = False + if not_in_axis: + out[out_i] = int64(data[i]) + out_i += 1 + + return out + + +# def _squeeze_shape_func_input_data(data, axis, ndims): +# out = [] +# for i in range(len(data)): +# if not i in axis: +# out.append(data[i]) +# return out + + +@_reg.register_shape_func("dyn.squeeze", [False, True]) +def dynamic_squeeze_shape_func(attrs, inputs, out_ndims): + return [_squeeze_shape_func_input_data(inputs[0], inputs[1], out_ndims[0])] \ No newline at end of file diff --git a/src/relay/op/dyn/tensor/transform.cc b/src/relay/op/dyn/tensor/transform.cc index d099432909c7..93172c542a6e 100644 --- a/src/relay/op/dyn/tensor/transform.cc +++ b/src/relay/op/dyn/tensor/transform.cc @@ -742,71 +742,6 @@ Expr MakeDynSqueeze(Expr data, Expr axes) { TVM_REGISTER_GLOBAL("relay.op.dyn._make.squeeze").set_body_typed(MakeDynSqueeze); - -// InferCorrectLayoutOutput SqueezeInferCorrectLayout(const Attrs& attrs, -// const Array& new_in_layouts, -// const Array& old_in_layouts, -// const Array& old_in_types) { -// const auto* attrs_ptr = attrs.as(); -// ICHECK(attrs_ptr); -// ObjectPtr params = make_object(*attrs_ptr); - -// Layout inferred_input = new_in_layouts.defined() ? new_in_layouts[0] : old_in_layouts[0]; -// Layout inferred_output = inferred_input; - -// ICHECK(old_in_types[0].as()); -// const auto& shape = old_in_types[0].as()->shape; - -// // axis to squeeze -// Array axis; -// if (params->axis.defined()) { -// axis = params->axis; -// } else { -// // if axes is None, squeeze all axes of dimension 1 -// for (size_t i = 0; i < shape.size(); i++) { -// if (topi::detail::GetConstInt(shape[i]) == 1) { -// axis.push_back(i); -// } -// } -// } - -// // If new_in_layouts are defined, this code tries to modify the layout -// if (new_in_layouts.defined() && old_in_layouts.defined()) { -// Array new_axis; -// for (const auto& e : axis) { -// const auto& dim = old_in_layouts[0][e]; -// new_axis.push_back((new_in_layouts[0]).IndexOf(dim)); -// } -// params->axis = new_axis; -// axis = new_axis; -// } - -// // Infer output layout -// Array kept_axes; -// for (size_t i = 0; i < inferred_input.ndim(); i++) { -// bool is_dim_kept = true; - -// // Check whether the dim should be kept -// for (const auto& e : axis) { -// int64_t axis_val = e->value; -// if (axis_val < 0) { -// axis_val += inferred_input.ndim(); -// } -// if (static_cast(i) == axis_val) { -// is_dim_kept = false; -// break; -// } -// } - -// if (is_dim_kept) { -// kept_axes.push_back(inferred_input->axes[i]); -// } -// } -// inferred_output = Layout(kept_axes); - -// return InferCorrectLayoutOutput({inferred_input}, {inferred_output}, Attrs(params)); -// } - RELAY_REGISTER_OP("dyn.squeeze") .describe(R"code(Squeeze the input tensor at the dimensions given by axes @@ -822,7 +757,6 @@ RELAY_REGISTER_OP("dyn.squeeze") .add_type_rel("DynSqueeze", DynSqueezeRel) .set_attr("FTVMCompute", SqueezeCompute) .set_attr("TOpPattern", kInjective) - // .set_attr("FInferCorrectLayout", SqueezeInferCorrectLayout) .set_attr("TReshapeOp", true); } // namespace dyn diff --git a/tests/python/frontend/onnx/test_forward.py b/tests/python/frontend/onnx/test_forward.py index e049a195dc17..072f128db038 100644 --- a/tests/python/frontend/onnx/test_forward.py +++ b/tests/python/frontend/onnx/test_forward.py @@ -4950,8 +4950,8 @@ def verify_eyelike(indata): "test_split_variable_parts_2d", "test_split_variable_parts_default_axis", "test_split_zero_size_splits", - "test_squeeze", - "test_squeeze_negative_axes", + # "test_squeeze", + # "test_squeeze_negative_axes", "test_strnormalizer_export_monday_casesensintive_lower", "test_strnormalizer_export_monday_casesensintive_nochangecase", "test_strnormalizer_export_monday_casesensintive_upper", From 4a234da4fe1f538d71673ccaf457b02547a35979 Mon Sep 17 00:00:00 2001 From: CircleSpin Date: Wed, 22 Sep 2021 16:17:48 -0400 Subject: [PATCH 04/11] add Squeeze op convert to onnx.py --- python/tvm/relay/frontend/onnx.py | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/python/tvm/relay/frontend/onnx.py b/python/tvm/relay/frontend/onnx.py index 49610d63ed8d..c0ee6f7a392b 100644 --- a/python/tvm/relay/frontend/onnx.py +++ b/python/tvm/relay/frontend/onnx.py @@ -1487,13 +1487,16 @@ class Squeeze(OnnxOpConverter): @classmethod def _impl_v1(cls, inputs, attr, params): + return _op.squeeze(*inputs, axis=attr["axes"]) @classmethod def _impl_v13(cls, inputs, attr, params): - rank = _op.shape_of(_op.shape_of(inputs[0])) axis = inputs[1] - axis = _op.where(axis < _op.const(0, *axis dtype*), axis + rank, axis) - return _op.squeeze(inputs[1], axis) + dtype = infer_type(axis).checked_type.dtype + rank = _op.shape_of(_op.shape_of(inputs[0], dtype), dtype) + axis = _op.where(axis < _op.const(0, dtype), axis + rank, axis) + return _op.squeeze(inputs[0], fold_constant(axis)) + class Split(OnnxOpConverter): """Operator converter for Split.""" @@ -2817,7 +2820,9 @@ def _impl_v12(cls, inputs, attr, params): alpha = _op.const(attr.get("alpha", 1.0), dtype) zero = _op.const(0, dtype) one = _op.const(1, dtype) - return _op.maximum(zero, x) + _op.minimum(zero, alpha * (_op.exp(x / alpha) - one)) + out = _op.maximum(zero, x) + _op.minimum(zero, alpha * (_op.exp(x / alpha) - one)) + print(out) + return out class MaxRoiPool(OnnxOpConverter): @@ -4147,7 +4152,7 @@ def _get_convert_map(opset): "ScatterElements": Scatter.get_converter(opset), "ScatterND": ScatterND.get_converter(opset), "EyeLike": EyeLike.get_converter(opset), - "Squeeze": AttrCvt("squeeze", {"axes": "axis"}), + "Squeeze": Squeeze.get_converter(opset), "Unsqueeze": Unsqueeze.get_converter(opset), "Pad": Pad.get_converter(opset), "Shape": Shape.get_converter(opset), From 7197c50efdc4f38d87cb3a761dbfc32616b76ac7 Mon Sep 17 00:00:00 2001 From: CircleSpin Date: Thu, 23 Sep 2021 13:09:26 -0400 Subject: [PATCH 05/11] dynamic to static --- src/relay/transforms/dynamic_to_static.cc | 9 ++++++++ tests/python/frontend/onnx/test_forward.py | 2 -- .../relay/test_pass_dynamic_to_static.py | 23 +++++++++++++++++++ 3 files changed, 32 insertions(+), 2 deletions(-) diff --git a/src/relay/transforms/dynamic_to_static.cc b/src/relay/transforms/dynamic_to_static.cc index 318022fb86f5..b4a7f96da455 100644 --- a/src/relay/transforms/dynamic_to_static.cc +++ b/src/relay/transforms/dynamic_to_static.cc @@ -45,6 +45,15 @@ class DynamicToStaticMutator : public MixedModeMutator { } return Expr(nullptr); }}, + {Op::Get("dyn.squeeze"), + [this](const CallNode* call_node) { + auto args = PrepareArgs(call_node); + if (const ConstantNode* axis = args[1].as()) { + ICHECK_EQ(axis->data->ndim, 1); + return MakeSqueeze(call_node->args[0], ToVector(axis->data)); + } + return Expr(nullptr); + }}, {Op::Get("dyn.tile"), [this](const CallNode* call_node) { auto args = PrepareArgs(call_node); diff --git a/tests/python/frontend/onnx/test_forward.py b/tests/python/frontend/onnx/test_forward.py index 072f128db038..a4ffcc39b52b 100644 --- a/tests/python/frontend/onnx/test_forward.py +++ b/tests/python/frontend/onnx/test_forward.py @@ -4950,8 +4950,6 @@ def verify_eyelike(indata): "test_split_variable_parts_2d", "test_split_variable_parts_default_axis", "test_split_zero_size_splits", - # "test_squeeze", - # "test_squeeze_negative_axes", "test_strnormalizer_export_monday_casesensintive_lower", "test_strnormalizer_export_monday_casesensintive_nochangecase", "test_strnormalizer_export_monday_casesensintive_upper", diff --git a/tests/python/relay/test_pass_dynamic_to_static.py b/tests/python/relay/test_pass_dynamic_to_static.py index a34c4ac6f705..ed8adb99e368 100644 --- a/tests/python/relay/test_pass_dynamic_to_static.py +++ b/tests/python/relay/test_pass_dynamic_to_static.py @@ -71,6 +71,29 @@ def verify_reshape(shape, newshape, oshape): verify_reshape((2, 3, 4), (8, 3), (8, 3)) verify_reshape((4, 7), (2, 7, 2), (2, 7, 2)) +@tvm.testing.uses_gpu +def test_dynamic_to_static_squeeze(): + def verify_squeeze(shape, axis, oshape): + x = relay.var("x", relay.TensorType(shape, "float32")) + y = relay.var("y", relay.TensorType(axis, "float32")) + z = relay.squeeze(x, relay.shape_of(y)) + func = run_infer_type(relay.Function([x, y], z)) + func2 = run_opt_pass(run_opt_pass(func, transform.DynamicToStatic()), transform.InferType()) + + zz = func2.body + assert isinstance(zz, relay.Call) + assert zz.op == relay.op.get("squeeze") + assert "axis=" in zz.astext() + assert zz.checked_type == relay.ty.TensorType(oshape, "float32") + + x_data = np.random.uniform(low=-1, high=1, size=shape).astype("float32") + y_data = np.random.uniform(low=-1, high=1, size=axis).astype("float32") + ref_res = np.squeeze(x_data, axis) + verify_func(func2, [x_data, y_data], ref_res) + + verify_squeeze((1, 3, 4, 1), (0,), (3, 4, 1)) + verify_squeeze((1, 3, 4, 1), (3,), (1, 3, 4)) + verify_squeeze((1, 3, 4, 1), (0, 3), (3, 4)) @tvm.testing.uses_gpu def test_dynamic_to_static_double_reshape(): From ee48312b1fb78b71d4bd3f260df680ec04db3118 Mon Sep 17 00:00:00 2001 From: CircleSpin Date: Thu, 23 Sep 2021 13:48:18 -0400 Subject: [PATCH 06/11] removed comments --- python/tvm/relay/op/dyn/_transform.py | 8 -------- src/relay/op/dyn/tensor/transform.cc | 13 +++---------- 2 files changed, 3 insertions(+), 18 deletions(-) diff --git a/python/tvm/relay/op/dyn/_transform.py b/python/tvm/relay/op/dyn/_transform.py index 1d855cb9420d..cbc6addda66a 100644 --- a/python/tvm/relay/op/dyn/_transform.py +++ b/python/tvm/relay/op/dyn/_transform.py @@ -277,14 +277,6 @@ def _squeeze_shape_func_input_data(data, axis, ndims): return out -# def _squeeze_shape_func_input_data(data, axis, ndims): -# out = [] -# for i in range(len(data)): -# if not i in axis: -# out.append(data[i]) -# return out - - @_reg.register_shape_func("dyn.squeeze", [False, True]) def dynamic_squeeze_shape_func(attrs, inputs, out_ndims): return [_squeeze_shape_func_input_data(inputs[0], inputs[1], out_ndims[0])] \ No newline at end of file diff --git a/src/relay/op/dyn/tensor/transform.cc b/src/relay/op/dyn/tensor/transform.cc index 93172c542a6e..ae208da6669f 100644 --- a/src/relay/op/dyn/tensor/transform.cc +++ b/src/relay/op/dyn/tensor/transform.cc @@ -724,14 +724,7 @@ Array SqueezeCompute(const Attrs& attrs, const Array& in newshape.push_back(val); } } - // newshape = {1, shape_var0} return {topi::reshape(inputs[0], newshape)}; - // def dyn_reshape(x, shape_var0): - // out = zeros((1 * shape_var0,)) - // for j in range(1): - // for k in range(shape_var0): - // out[j*shape_var0 + k] = x[j * shape_var0 + k] - // return out } Expr MakeDynSqueeze(Expr data, Expr axes) { @@ -759,6 +752,6 @@ RELAY_REGISTER_OP("dyn.squeeze") .set_attr("TOpPattern", kInjective) .set_attr("TReshapeOp", true); -} // namespace dyn -} // namespace relay -} // namespace tvm +} +} +} From 150478288990e187703a4260105deb89e2406760 Mon Sep 17 00:00:00 2001 From: CircleSpin Date: Fri, 24 Sep 2021 21:26:26 -0400 Subject: [PATCH 07/11] removed comments --- python/tvm/relay/frontend/onnx.py | 4 ++-- python/tvm/relay/op/transform.py | 2 +- src/relay/op/dyn/tensor/transform.cc | 16 +++++++--------- tests/python/relay/dyn/test_dynamic_op_level3.py | 4 ---- 4 files changed, 10 insertions(+), 16 deletions(-) diff --git a/python/tvm/relay/frontend/onnx.py b/python/tvm/relay/frontend/onnx.py index c0ee6f7a392b..a8456e0df047 100644 --- a/python/tvm/relay/frontend/onnx.py +++ b/python/tvm/relay/frontend/onnx.py @@ -1487,7 +1487,8 @@ class Squeeze(OnnxOpConverter): @classmethod def _impl_v1(cls, inputs, attr, params): - return _op.squeeze(*inputs, axis=attr["axes"]) + axis = attr.get("axes", None) + return _op.squeeze(*inputs, axis) @classmethod def _impl_v13(cls, inputs, attr, params): @@ -2821,7 +2822,6 @@ def _impl_v12(cls, inputs, attr, params): zero = _op.const(0, dtype) one = _op.const(1, dtype) out = _op.maximum(zero, x) + _op.minimum(zero, alpha * (_op.exp(x / alpha) - one)) - print(out) return out diff --git a/python/tvm/relay/op/transform.py b/python/tvm/relay/op/transform.py index 234e76b11813..d0a3768a455f 100644 --- a/python/tvm/relay/op/transform.py +++ b/python/tvm/relay/op/transform.py @@ -1765,4 +1765,4 @@ def invert_permutation(data): data = [3, 4, 0, 2, 1] relay.invert_permutation(data) = [2, 4, 3, 0, 1] """ - return _make.invert_permutation(data) + return _make.invert_permutation(data) \ No newline at end of file diff --git a/src/relay/op/dyn/tensor/transform.cc b/src/relay/op/dyn/tensor/transform.cc index ae208da6669f..d8f6457e8593 100644 --- a/src/relay/op/dyn/tensor/transform.cc +++ b/src/relay/op/dyn/tensor/transform.cc @@ -694,6 +694,7 @@ RELAY_REGISTER_OP("dyn.expand_dims") bool DynSqueezeRel(const Array& types, int num_inputs, const Attrs& attrs, const TypeReporter& reporter) { + // [data, axes, output] ICHECK_EQ(types.size(), 3); const auto* data = types[0].as(); if (data == nullptr) { @@ -703,7 +704,8 @@ bool DynSqueezeRel(const Array& types, int num_inputs, const Attrs& attrs, if (axes == nullptr) { return false; } - ICHECK_EQ(axes->shape.size(), 1); + + ICHECK_EQ(axes->shape.size(), 1) << "Got" << axes->shape.size() << "expected 1"; ICHECK(axes->shape[0].as()); size_t output_rank = data->shape.size() - axes->shape[0].as()->value; std::vector result_shape(output_rank, Any()); @@ -718,11 +720,7 @@ Array SqueezeCompute(const Attrs& attrs, const Array& in ICHECK(out_ttype != nullptr); Array newshape; for (auto val : out_ttype->shape) { - if (val->IsInstance()) { newshape.push_back(val.as()->ToVar()); - } else { - newshape.push_back(val); - } } return {topi::reshape(inputs[0], newshape)}; } @@ -736,7 +734,7 @@ Expr MakeDynSqueeze(Expr data, Expr axes) { TVM_REGISTER_GLOBAL("relay.op.dyn._make.squeeze").set_body_typed(MakeDynSqueeze); RELAY_REGISTER_OP("dyn.squeeze") - .describe(R"code(Squeeze the input tensor at the dimensions given by axes + .describe(R"code(Remove axes of value 1 in input tensor at the dimensions given by axes - **data**: The input data to the operator. - **axes**: The axes to squeeze. @@ -752,6 +750,6 @@ RELAY_REGISTER_OP("dyn.squeeze") .set_attr("TOpPattern", kInjective) .set_attr("TReshapeOp", true); -} -} -} +} // namespace dyn +} // namespace dyn +} // namespace tvm diff --git a/tests/python/relay/dyn/test_dynamic_op_level3.py b/tests/python/relay/dyn/test_dynamic_op_level3.py index f33743b9be02..259a936a88e7 100644 --- a/tests/python/relay/dyn/test_dynamic_op_level3.py +++ b/tests/python/relay/dyn/test_dynamic_op_level3.py @@ -99,11 +99,7 @@ def verify_squeeze(shape, dtype, axis): np_axis = tuple(axis) axis = relay.var("axis", relay.TensorType([len(axis)], "int64")) squeeze = relay.squeeze(x, axis=axis) - - print(squeeze) func = relay.Function([x, axis], squeeze) - print(func) - x_data = np.random.random_sample(shape).astype(dtype) ref_res = np.squeeze(x_data, axis=np_axis) verify_func(func, [x_data, np.array(np_axis).astype("int64")], ref_res) From ef4796975d4a1712a8523a82faca7f9b027404eb Mon Sep 17 00:00:00 2001 From: CircleSpin Date: Fri, 24 Sep 2021 21:26:57 -0400 Subject: [PATCH 08/11] added comment --- src/relay/op/dyn/tensor/transform.cc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/relay/op/dyn/tensor/transform.cc b/src/relay/op/dyn/tensor/transform.cc index d8f6457e8593..8adfee595ca8 100644 --- a/src/relay/op/dyn/tensor/transform.cc +++ b/src/relay/op/dyn/tensor/transform.cc @@ -706,7 +706,7 @@ bool DynSqueezeRel(const Array& types, int num_inputs, const Attrs& attrs, } ICHECK_EQ(axes->shape.size(), 1) << "Got" << axes->shape.size() << "expected 1"; - ICHECK(axes->shape[0].as()); + ICHECK(axes->shape[0].as()) << "axes expected to be static"; size_t output_rank = data->shape.size() - axes->shape[0].as()->value; std::vector result_shape(output_rank, Any()); reporter->Assign(types[2], TensorType(result_shape, data->dtype)); From 46a5f57dd0f225447bdc7300dbdf35977e7637a4 Mon Sep 17 00:00:00 2001 From: CircleSpin Date: Fri, 24 Sep 2021 21:32:08 -0400 Subject: [PATCH 09/11] adjusted comment --- src/relay/op/dyn/tensor/transform.cc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/relay/op/dyn/tensor/transform.cc b/src/relay/op/dyn/tensor/transform.cc index 8adfee595ca8..c548d071f46e 100644 --- a/src/relay/op/dyn/tensor/transform.cc +++ b/src/relay/op/dyn/tensor/transform.cc @@ -706,7 +706,7 @@ bool DynSqueezeRel(const Array& types, int num_inputs, const Attrs& attrs, } ICHECK_EQ(axes->shape.size(), 1) << "Got" << axes->shape.size() << "expected 1"; - ICHECK(axes->shape[0].as()) << "axes expected to be static"; + ICHECK(axes->shape[0].as()) << "axes expected to be static rank"; size_t output_rank = data->shape.size() - axes->shape[0].as()->value; std::vector result_shape(output_rank, Any()); reporter->Assign(types[2], TensorType(result_shape, data->dtype)); From 95996d4b200cc79be374b74b59f6f59aa85a12e4 Mon Sep 17 00:00:00 2001 From: CircleSpin Date: Fri, 24 Sep 2021 21:44:02 -0400 Subject: [PATCH 10/11] black and lint --- python/tvm/relay/frontend/onnx.py | 2 ++ python/tvm/relay/op/dyn/_transform.py | 2 +- python/tvm/relay/op/transform.py | 2 +- src/relay/op/dyn/tensor/transform.cc | 2 +- tests/python/relay/dyn/test_dynamic_op_level3.py | 1 + tests/python/relay/test_pass_dynamic_to_static.py | 2 ++ 6 files changed, 8 insertions(+), 3 deletions(-) diff --git a/python/tvm/relay/frontend/onnx.py b/python/tvm/relay/frontend/onnx.py index a8456e0df047..0f99193fe525 100644 --- a/python/tvm/relay/frontend/onnx.py +++ b/python/tvm/relay/frontend/onnx.py @@ -1482,6 +1482,7 @@ def _impl_v12(cls, inputs, attr, params): result = _op.expand_dims(result, axis) return result + class Squeeze(OnnxOpConverter): """Operator converter for Squeeze.""" @@ -1498,6 +1499,7 @@ def _impl_v13(cls, inputs, attr, params): axis = _op.where(axis < _op.const(0, dtype), axis + rank, axis) return _op.squeeze(inputs[0], fold_constant(axis)) + class Split(OnnxOpConverter): """Operator converter for Split.""" diff --git a/python/tvm/relay/op/dyn/_transform.py b/python/tvm/relay/op/dyn/_transform.py index cbc6addda66a..c909764319d9 100644 --- a/python/tvm/relay/op/dyn/_transform.py +++ b/python/tvm/relay/op/dyn/_transform.py @@ -279,4 +279,4 @@ def _squeeze_shape_func_input_data(data, axis, ndims): @_reg.register_shape_func("dyn.squeeze", [False, True]) def dynamic_squeeze_shape_func(attrs, inputs, out_ndims): - return [_squeeze_shape_func_input_data(inputs[0], inputs[1], out_ndims[0])] \ No newline at end of file + return [_squeeze_shape_func_input_data(inputs[0], inputs[1], out_ndims[0])] diff --git a/python/tvm/relay/op/transform.py b/python/tvm/relay/op/transform.py index d0a3768a455f..234e76b11813 100644 --- a/python/tvm/relay/op/transform.py +++ b/python/tvm/relay/op/transform.py @@ -1765,4 +1765,4 @@ def invert_permutation(data): data = [3, 4, 0, 2, 1] relay.invert_permutation(data) = [2, 4, 3, 0, 1] """ - return _make.invert_permutation(data) \ No newline at end of file + return _make.invert_permutation(data) diff --git a/src/relay/op/dyn/tensor/transform.cc b/src/relay/op/dyn/tensor/transform.cc index c548d071f46e..06a0ce3da33f 100644 --- a/src/relay/op/dyn/tensor/transform.cc +++ b/src/relay/op/dyn/tensor/transform.cc @@ -751,5 +751,5 @@ RELAY_REGISTER_OP("dyn.squeeze") .set_attr("TReshapeOp", true); } // namespace dyn -} // namespace dyn +} // namespace relay } // namespace tvm diff --git a/tests/python/relay/dyn/test_dynamic_op_level3.py b/tests/python/relay/dyn/test_dynamic_op_level3.py index 259a936a88e7..22583eda4a40 100644 --- a/tests/python/relay/dyn/test_dynamic_op_level3.py +++ b/tests/python/relay/dyn/test_dynamic_op_level3.py @@ -107,6 +107,7 @@ def verify_squeeze(shape, dtype, axis): verify_squeeze((1, 3, 1), "float32", [0]) verify_squeeze((1, 2, 1, 2, 1), "float32", [0, 2]) + @tvm.testing.uses_gpu def test_dyn_expand_dims(): def verify_expand_dims( diff --git a/tests/python/relay/test_pass_dynamic_to_static.py b/tests/python/relay/test_pass_dynamic_to_static.py index ed8adb99e368..5b61733bbd76 100644 --- a/tests/python/relay/test_pass_dynamic_to_static.py +++ b/tests/python/relay/test_pass_dynamic_to_static.py @@ -71,6 +71,7 @@ def verify_reshape(shape, newshape, oshape): verify_reshape((2, 3, 4), (8, 3), (8, 3)) verify_reshape((4, 7), (2, 7, 2), (2, 7, 2)) + @tvm.testing.uses_gpu def test_dynamic_to_static_squeeze(): def verify_squeeze(shape, axis, oshape): @@ -95,6 +96,7 @@ def verify_squeeze(shape, axis, oshape): verify_squeeze((1, 3, 4, 1), (3,), (1, 3, 4)) verify_squeeze((1, 3, 4, 1), (0, 3), (3, 4)) + @tvm.testing.uses_gpu def test_dynamic_to_static_double_reshape(): def verify_reshape(shape, newshape): From e20f8a817cf59e9f593a0f5a5d22587649d8f8fc Mon Sep 17 00:00:00 2001 From: CircleSpin Date: Mon, 27 Sep 2021 13:17:23 -0400 Subject: [PATCH 11/11] ran make format in root directory --- src/relay/op/dyn/tensor/transform.cc | 5 ++--- src/relay/transforms/dynamic_to_static.cc | 2 +- 2 files changed, 3 insertions(+), 4 deletions(-) diff --git a/src/relay/op/dyn/tensor/transform.cc b/src/relay/op/dyn/tensor/transform.cc index 06a0ce3da33f..64baa6066522 100644 --- a/src/relay/op/dyn/tensor/transform.cc +++ b/src/relay/op/dyn/tensor/transform.cc @@ -693,7 +693,7 @@ RELAY_REGISTER_OP("dyn.expand_dims") .set_attr("TOpPattern", kInjective); bool DynSqueezeRel(const Array& types, int num_inputs, const Attrs& attrs, - const TypeReporter& reporter) { + const TypeReporter& reporter) { // [data, axes, output] ICHECK_EQ(types.size(), 3); const auto* data = types[0].as(); @@ -713,14 +713,13 @@ bool DynSqueezeRel(const Array& types, int num_inputs, const Attrs& attrs, return true; } - Array SqueezeCompute(const Attrs& attrs, const Array& inputs, const Type& out_type) { const auto* out_ttype = out_type.as(); ICHECK(out_ttype != nullptr); Array newshape; for (auto val : out_ttype->shape) { - newshape.push_back(val.as()->ToVar()); + newshape.push_back(val.as()->ToVar()); } return {topi::reshape(inputs[0], newshape)}; } diff --git a/src/relay/transforms/dynamic_to_static.cc b/src/relay/transforms/dynamic_to_static.cc index b4a7f96da455..751271d2add3 100644 --- a/src/relay/transforms/dynamic_to_static.cc +++ b/src/relay/transforms/dynamic_to_static.cc @@ -45,7 +45,7 @@ class DynamicToStaticMutator : public MixedModeMutator { } return Expr(nullptr); }}, - {Op::Get("dyn.squeeze"), + {Op::Get("dyn.squeeze"), [this](const CallNode* call_node) { auto args = PrepareArgs(call_node); if (const ConstantNode* axis = args[1].as()) {