diff --git a/mindspore/ccsrc/kernel/tbe/tbe_adapter.cc b/mindspore/ccsrc/kernel/tbe/tbe_adapter.cc index 005c290abab7a18a822020deaf11db9f5e87c391..11d2185c32ede7996352bc245dd53797ae826363 100644 --- a/mindspore/ccsrc/kernel/tbe/tbe_adapter.cc +++ b/mindspore/ccsrc/kernel/tbe/tbe_adapter.cc @@ -75,6 +75,8 @@ static std::map tbe_func_adapter_map = { {"resize_nearest_neighbor", "resize_nearest_neighbor_v2_d"}, {"resize_nearest_neighbor_grad", "resize_nearest_neighbor_v2_grad_d"}, {"pad", "pad_d"}, + {"space_to_batch", "space_to_batch_d"}, + {"batch_to_space", "batch_to_space_d"}, {"adam", "apply_adam_d"}}; void TbeAdapter::NormalizeFuncName(std::string *func_name) { diff --git a/mindspore/ops/_op_impl/tbe/__init__.py b/mindspore/ops/_op_impl/tbe/__init__.py index 738c9ef4723e3b6b5a8f56ab39510fc0c8cede8c..cdacb22a282e3b7df485092c689d042c98063b49 100644 --- a/mindspore/ops/_op_impl/tbe/__init__.py +++ b/mindspore/ops/_op_impl/tbe/__init__.py @@ -156,3 +156,5 @@ from .scatter_nd_update import _scatter_nd_update_tbe from .avg_pool import _avg_pool_tbe from .avg_pool_grad import _avg_pool_grad_tbe from .ones_like import _ones_like_tbe +from .batch_to_space import _batch_to_space_tbe +from .space_to_batch import _space_to_batch_tbe diff --git a/mindspore/ops/_op_impl/tbe/batch_to_space.py b/mindspore/ops/_op_impl/tbe/batch_to_space.py new file mode 100644 index 0000000000000000000000000000000000000000..c327419284a02fc33a1c863c406bd3f7731c251b --- /dev/null +++ b/mindspore/ops/_op_impl/tbe/batch_to_space.py @@ -0,0 +1,38 @@ +# Copyright 2020 Huawei Technologies Co., Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================ + +"""BatchToSpace op""" +from mindspore.ops.op_info_register import op_info_register, TBERegOp, DataType + +batch_to_space_op_info = TBERegOp("BatchToSpace") \ + .fusion_type("OPAQUE") \ + .async_flag(False) \ + .binfile_name("batch_to_space_d.so") \ + .compute_cost(10) \ + .kernel_name("batch_to_space_d") \ + .partial_flag(True) \ + .attr("block_size", "required", "int", "all") \ + .attr("crops", "required", "listListInt", "all") \ + .input(0, "x", False, "required", "all") \ + .output(0, "y", False, "required", "all") \ + .dtype_format(DataType.F16_5HD, DataType.F16_5HD) \ + .dtype_format(DataType.F32_5HD, DataType.F32_5HD) \ + .get_op_info() + + +@op_info_register(batch_to_space_op_info) +def _batch_to_space_tbe(): + """BatchToSpace TBE register""" + return diff --git a/mindspore/ops/_op_impl/tbe/space_to_batch.py b/mindspore/ops/_op_impl/tbe/space_to_batch.py new file mode 100644 index 0000000000000000000000000000000000000000..2a49077899b90ee379cea57f4793f26ce0d59f0a --- /dev/null +++ b/mindspore/ops/_op_impl/tbe/space_to_batch.py @@ -0,0 +1,38 @@ +# Copyright 2020 Huawei Technologies Co., Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================ + +"""SpaceToBatch op""" +from mindspore.ops.op_info_register import op_info_register, TBERegOp, DataType + +space_to_batch_op_info = TBERegOp("SpaceToBatch") \ + .fusion_type("OPAQUE") \ + .async_flag(False) \ + .binfile_name("space_to_batch_d.so") \ + .compute_cost(10) \ + .kernel_name("space_to_batch_d") \ + .partial_flag(True) \ + .attr("block_size", "required", "int", "all") \ + .attr("paddings", "required", "listListInt", "all") \ + .input(0, "x", False, "required", "all") \ + .output(0, "y", False, "required", "all") \ + .dtype_format(DataType.F16_5HD, DataType.F16_5HD) \ + .dtype_format(DataType.F32_5HD, DataType.F32_5HD) \ + .get_op_info() + + +@op_info_register(space_to_batch_op_info) +def _space_to_batch_tbe(): + """SpaceToBatch TBE register""" + return diff --git a/tests/ut/python/ops/test_array_ops.py b/tests/ut/python/ops/test_array_ops.py index 61b8d48feafdf7c616952ddc2610723914e541e8..98945254261803c5e030f3086821e713d6117f06 100644 --- a/tests/ut/python/ops/test_array_ops.py +++ b/tests/ut/python/ops/test_array_ops.py @@ -95,6 +95,7 @@ def test_select(): expect = np.array([[1, 8, 9], [10, 5, 6]]) assert np.all(output.asnumpy() == expect) + def test_argmin_invalid_output_type(): P.Argmin(-1, mstype.int64) P.Argmin(-1, mstype.int32) @@ -203,6 +204,28 @@ class MathBinaryNet2(Cell): return self.logic_or(ret_less_equal, ret_greater) +class BatchToSpaceNet(Cell): + def __init__(self): + super(BatchToSpaceNet, self).__init__() + block_size = 2 + crops = [[0, 0], [0, 0]] + self.batch_to_space = P.BatchToSpace(block_size, crops) + + def construct(self, x): + return self.batch_to_space(x) + + +class SpaceToBatchNet(Cell): + def __init__(self): + super(SpaceToBatchNet, self).__init__() + block_size = 2 + paddings = [[0, 0], [0, 0]] + self.space_to_batch = P.SpaceToBatch(block_size, paddings) + + def construct(self, x): + return self.space_to_batch(x) + + test_case_array_ops = [ ('CustNet1', { 'block': CustNet1(), @@ -219,6 +242,12 @@ test_case_array_ops = [ ('MathBinaryNet2', { 'block': MathBinaryNet2(), 'desc_inputs': [Tensor(np.ones([2, 2]), dtype=ms.int32)]}), + ('BatchToSpaceNet', { + 'block': BatchToSpaceNet(), + 'desc_inputs': [Tensor(np.array([[[[1]]], [[[2]]], [[[3]]], [[[4]]]]).astype(np.float16))]}), + ('SpaceToBatchNet', { + 'block': SpaceToBatchNet(), + 'desc_inputs': [Tensor(np.array([[[[1, 2], [3, 4]]]]).astype(np.float16))]}), ] test_case_lists = [test_case_array_ops]