diff --git a/mindspore/ccsrc/transform/convert.cc b/mindspore/ccsrc/transform/convert.cc index c8c6abea4c9..c975b18812e 100755 --- a/mindspore/ccsrc/transform/convert.cc +++ b/mindspore/ccsrc/transform/convert.cc @@ -180,6 +180,8 @@ const char kNamePrint[] = "Print"; const char kNameApplyFtrl[] = "ApplyFtrl"; const char kNameDiag[] = "Diag"; const char kNameDiagPart[] = "DiagPart"; +const char kNameSpaceToBatch[] = "SpaceToBatch"; +const char kNameBatchToSpace[] = "BatchToSpace"; // -----------------OpAdapter initialization-------------- std::unordered_map &DfGraphConvertor::get_adpt_map() { @@ -361,7 +363,9 @@ std::unordered_map &DfGraphConvertor::get_adpt_ma {string(kNameRound), ADPT_DESC(Round)}, {string(kNameApplyFtrl), ADPT_DESC(ApplyFtrl)}, {string(kNameDiag), ADPT_DESC(Diag)}, - {string(kNameDiagPart), ADPT_DESC(DiagPart)}}; + {string(kNameDiagPart), ADPT_DESC(DiagPart)}, + {string(kNameSpaceToBatch), ADPT_DESC(SpaceToBatchD)}, + {string(kNameBatchToSpace), ADPT_DESC(BatchToSpaceD)}}; #ifdef ENABLE_GE adpt_map[string(kNamePrint)] = ADPT_DESC(Print); #endif diff --git a/mindspore/ccsrc/transform/op_adapter.h b/mindspore/ccsrc/transform/op_adapter.h index aa466adbb8a..3dd299f83de 100644 --- a/mindspore/ccsrc/transform/op_adapter.h +++ b/mindspore/ccsrc/transform/op_adapter.h @@ -744,6 +744,28 @@ class OpAdapter : public BaseOpAdapter { return list; } + static std::vector ConvertAny(const ValuePtr& value, const AnyTraits>>, + const AnyTraits>) { + MS_EXCEPTION_IF_NULL(value); + MS_LOG(DEBUG) << "Value: " << value->type_name(); + if (!value->isa()) { + MS_LOG(EXCEPTION) << "Value should be ValueList, but got " << value->type_name(); + } + auto vec = value->cast(); + std::vector list; + for (auto& it : vec->value()) { + MS_EXCEPTION_IF_NULL(it); + if (!it->isa()) { + MS_LOG(EXCEPTION) << "It should be ValueList, but got " << it->type_name(); + } + auto sub_vector = it->cast(); + for (auto& item : sub_vector->value()) { + list.push_back(static_cast(GetValue(item))); + } + } + return list; + } + static std::vector ConvertAny(const ValuePtr& value, const AnyTraits>, const AnyTraits>) { MS_EXCEPTION_IF_NULL(value); diff --git a/mindspore/ccsrc/transform/op_declare.cc b/mindspore/ccsrc/transform/op_declare.cc index af0ce29f859..cc63bacd8a5 100755 --- a/mindspore/ccsrc/transform/op_declare.cc +++ b/mindspore/ccsrc/transform/op_declare.cc @@ -1183,6 +1183,19 @@ INPUT_MAP(DiagPart) = {{1, INPUT_DESC(x)}}; ATTR_MAP(DiagPart) = EMPTY_ATTR_MAP; OUTPUT_MAP(DiagPart) = {{0, OUTPUT_DESC(y)}}; +// SpaceToBatchD +INPUT_MAP(SpaceToBatchD) = {{1, INPUT_DESC(x)}}; +ATTR_MAP(SpaceToBatchD) = { + {"block_size", ATTR_DESC(block_size, AnyTraits())}, + {"paddings", ATTR_DESC(paddings, AnyTraits>>(), AnyTraits>())}}; +OUTPUT_MAP(SpaceToBatchD) = {{0, OUTPUT_DESC(y)}}; + +// BatchToSpaceD +INPUT_MAP(BatchToSpaceD) = {{1, INPUT_DESC(x)}}; +ATTR_MAP(BatchToSpaceD) = { + {"block_size", ATTR_DESC(block_size, AnyTraits())}, + {"crops", ATTR_DESC(crops, AnyTraits>>(), AnyTraits>())}}; +OUTPUT_MAP(BatchToSpaceD) = {{0, OUTPUT_DESC(y)}}; #ifdef ENABLE_GE // Print INPUT_MAP(Print) = EMPTY_INPUT_MAP; diff --git a/mindspore/ccsrc/transform/op_declare.h b/mindspore/ccsrc/transform/op_declare.h index 59014b8b7d7..978828e16d1 100755 --- a/mindspore/ccsrc/transform/op_declare.h +++ b/mindspore/ccsrc/transform/op_declare.h @@ -439,6 +439,10 @@ DECLARE_OP_ADAPTER(Diag) DECLARE_OP_USE_OUTPUT(Diag) DECLARE_OP_ADAPTER(DiagPart) DECLARE_OP_USE_OUTPUT(DiagPart) +DECLARE_OP_ADAPTER(SpaceToBatchD) +DECLARE_OP_USE_OUTPUT(SpaceToBatchD) +DECLARE_OP_ADAPTER(BatchToSpaceD) +DECLARE_OP_USE_OUTPUT(BatchToSpaceD) #ifdef ENABLE_GE DECLARE_OP_ADAPTER(Print) DECLARE_OP_USE_DYN_INPUT(Print) diff --git a/mindspore/ops/_grad/grad_array_ops.py b/mindspore/ops/_grad/grad_array_ops.py index 79841cf27aa..81d38a1e1e3 100644 --- a/mindspore/ops/_grad/grad_array_ops.py +++ b/mindspore/ops/_grad/grad_array_ops.py @@ -430,3 +430,23 @@ def get_bprop_diag_part(self): return (op(dout),) return bprop + + +@bprop_getters.register(P.SpaceToBatch) +def get_bprop_space_to_batch(self): + """Generate bprop for SpaceToBatch""" + space_to_batch_grad = P.BatchToSpace(self.block_size, self.paddings) + def bprop(x, out, dout): + dx = space_to_batch_grad(dout) + return (dx,) + return bprop + + +@bprop_getters.register(P.BatchToSpace) +def get_bprop_batch_to_space(self): + """Generate bprop for BatchToSpace""" + batch_to_space_grad = P.SpaceToBatch(self.block_size, self.crops) + def bprop(x, out, dout): + dx = batch_to_space_grad(dout) + return (dx,) + return bprop diff --git a/mindspore/ops/operations/__init__.py b/mindspore/ops/operations/__init__.py index 295e2edaf8a..899b2f8d0f5 100644 --- a/mindspore/ops/operations/__init__.py +++ b/mindspore/ops/operations/__init__.py @@ -29,7 +29,7 @@ from .array_ops import (Argmax, Argmin, Cast, ConcatOffset, Concat, Shape, Size, Slice, Split, Squeeze, StridedSlice, Tile, Transpose, TruncatedNormal, TupleToArray, - UnsortedSegmentSum, SpaceToDepth, DepthToSpace) + UnsortedSegmentSum, SpaceToDepth, DepthToSpace, SpaceToBatch, BatchToSpace) from .comm_ops import (AllGather, AllReduce, _AlltoAll, ReduceScatter, Broadcast, _MirrorOperator, ReduceOp, _VirtualDataset, _VirtualDiv, _GetTensorSlice) @@ -225,6 +225,8 @@ __all__ = [ "LARSUpdate", "Round", "ApplyFtrl", + "SpaceToBatch", + "BatchToSpace" ] __all__.sort() diff --git a/mindspore/ops/operations/array_ops.py b/mindspore/ops/operations/array_ops.py index 6740f172b45..8585f873e98 100644 --- a/mindspore/ops/operations/array_ops.py +++ b/mindspore/ops/operations/array_ops.py @@ -20,6 +20,7 @@ import copy import functools +import itertools import numbers import numpy as np @@ -2020,3 +2021,143 @@ class DepthToSpace(PrimitiveWithInfer): def infer_dtype(self, x_dtype): validator.check_subclass("x_dtype", x_dtype, mstype.tensor) return x_dtype + + +class SpaceToBatch(PrimitiveWithInfer): + r""" + Divide spatial dimensions into blocks and combine the block size with the original batch. + + This operation will divide spatial dimensions (H, W) into blocks with block_size, the output tensor's H and W + dimension is the corresponding number of blocks after division. The output tensor's batch dimension is the + product of the original batch and the square of block_size. Prior to division into blocks, the spatial dimensions + of the input are zero padded according to paddings if necessary. + + Args: + block_size (int): The block size of dividing block with value >= 1. + paddings (list): The padding value for H and W dimension, containing 2 sub list, each containing 2 int value. + All values must be >= 0. paddings[i] specifies the paddings for spatial dimension i, which corresponds to + input dimension i+2. It is required that input_shape[i+2]+paddings[i][0]+paddings[i][1] is divisible + by block_size. + + Inputs: + - **input_x** (Tensor) - The input tensor. + + Outputs: + Tensor, the output tensor with the same type as input. Assume input shape is :math:`(n, c, h, w)` with + :math:`block\_size` and :math:`padddings`. The output tensor shape will be :math:`(n', c', h', w')`, where + + :math:`n' = n*(block\_size*block\_size)` + + :math:`c' = c` + + :math:`h' = (h+paddings[0][0]+paddings[0][1])//block\_size` + + :math:`w' = (w+paddings[1][0]+paddings[1][1])//block\_size` + + Examples: + >>> block_size = 2 + >>> paddings = [[0, 0], [0, 0]] + >>> space_to_batch = P.SpaceToBatch(block_size, paddings) + >>> x = Tensor(np.array([[[[1, 2], [3, 4]]]]), mstype.float32) + >>> space_to_batch(x) + [[[[1.]]], [[[2.]]], [[[3.]]], [[[4.]]]] + + """ + @prim_attr_register + def __init__(self, block_size, paddings): + """Init SpaceToBatch""" + validator.check_type('block_size', block_size, [int]) + validator.check('block_size', block_size, '', 1, Rel.GT) + self.block_size = block_size + validator.check('paddings shape', np.array(paddings).shape, '', (2, 2)) + for elem in itertools.chain(*paddings): + validator.check_type('paddings element', elem, [int]) + self.paddings = paddings + + def infer_dtype(self, x_dtype): + validator.check_subclass("input_x", x_dtype, mstype.tensor) + validator.check_typename('input_x', x_dtype, mstype.number_type) + return x_dtype + + def infer_shape(self, x_shape): + validator.check('rank of input_x', len(x_shape), '', 4) + out_shape = copy.deepcopy(x_shape) + for i in range(2): + padded = out_shape[i+2] + self.paddings[i][0] + \ + self.paddings[i][1] + if padded % self.block_size != 0: + raise ValueError(f'padded[{i}] {padded} should be divisible by ' + f'block_size {self.block_size}') + out_shape[i+2] = padded // self.block_size + out_shape[0] *= self.block_size * self.block_size + return out_shape + + +class BatchToSpace(PrimitiveWithInfer): + r""" + Divide batch dimension with blocks and interleaves these blocks back into spatial dimensions. + + This operation will divide batch dimension N into blocks with block_size, the output tensor's N dimension + is the corresponding number of blocks after division. The output tensor's H, W dimension is product of original H, W + dimension and block_size with given amount to crop from dimension, respectively. + + Args: + block_size (int): The block size of dividing block with value >= 1. + crops (list): The crop value for H and W dimension, containing 2 sub list, each containing 2 int value. + All values must be >= 0. crops[i] specifies the crop values for spatial dimension i, which corresponds to + input dimension i+2. It is required that input_shape[i+2]*block_size >= crops[i][0]+crops[i][1]. + + Inputs: + - **input_x** (Tensor) - The input tensor. + + Outputs: + Tensor, the output tensor with the same type as input. Assume input shape is (n, c, h, w) with block_size + and crops. The output shape will be (n', c', h', w'), where + + :math:`n' = n//(block\_size*block\_size)` + + :math:`c' = c` + + :math:`h' = h*block\_size-crops[0][0]-crops[0][1]` + + :math:`w' = w*block\_size-crops[1][0]-crops[1][1]` + + Examples: + >>> block_size = 2 + >>> crops = [[0, 0], [0, 0]] + >>> op = P.BatchToSpace(block_size, crops) + >>> x = Tensor(np.array([[[[1]]], [[[2]]], [[[3]]], [[[4]]]]), mstype.float32) + >>> output = op(x) + [[[[1., 2.], [3., 4.]]]] + + """ + @prim_attr_register + def __init__(self, block_size, crops): + """Init BatchToSpace""" + validator.check_type('block_size', block_size, [int]) + validator.check('block_size', block_size, '', 1, Rel.GT) + self.block_size = block_size + validator.check('crops shape', np.array(crops).shape, '', (2, 2)) + for elem in itertools.chain(*crops): + validator.check_type('crops element', elem, [int]) + self.crops = crops + + def infer_dtype(self, x_dtype): + validator.check_subclass("input_x", x_dtype, mstype.tensor) + validator.check_typename('input_x', x_dtype, mstype.number_type) + return x_dtype + + def infer_shape(self, x_shape): + validator.check('rank of input_x', len(x_shape), '', 4) + out_shape = copy.deepcopy(x_shape) + for i in range(2): + x_block_prod = out_shape[i+2] * self.block_size + crops_sum = self.crops[i][0] + self.crops[i][1] + validator.check("x block shape prod", x_block_prod, 'crops sum', crops_sum, Rel.GT) + out_shape[i+2] = x_block_prod - crops_sum + block_size_prod = self.block_size * self.block_size + if out_shape[0] % block_size_prod != 0: + raise ValueError(f'input_x dimension 0 {out_shape[0]} should be divisible by ' + f'block_size_prod {block_size_prod}') + out_shape[0] = out_shape[0] // block_size_prod + return out_shape diff --git a/tests/ut/python/ops/test_ops.py b/tests/ut/python/ops/test_ops.py index e917c12748b..f1d365e9cf6 100755 --- a/tests/ut/python/ops/test_ops.py +++ b/tests/ut/python/ops/test_ops.py @@ -952,6 +952,26 @@ test_case_array_ops = [ 'desc_inputs': [[4, 4]], 'desc_bprop': [[4]], }), + ('SpaceToBatch_1', { + 'block': P.SpaceToBatch(2, [[0, 0], [0, 0]]), + 'desc_inputs': [[1, 3, 2, 2]], + 'desc_bprop': [[4, 3, 1, 1]], + }), + ('SpaceToBatch_2', { + 'block': P.SpaceToBatch(2, [[1, 1], [0, 4]]), + 'desc_inputs': [[1, 3, 2, 2]], + 'desc_bprop': [[4, 3, 2, 4]], + }), + ('BatchToSpace_1', { + 'block': P.BatchToSpace(2, [[0, 0], [0, 0]]), + 'desc_inputs': [[4, 3, 1, 1]], + 'desc_bprop': [[1, 3, 2, 2]], + }), + ('BatchToSpace_2', { + 'block': P.BatchToSpace(2, [[0, 0], [0, 1]]), + 'desc_inputs': [[4, 3, 1, 1]], + 'desc_bprop': [[1, 3, 2, 1]], + }), ] test_case_other_ops = [