diff --git a/mindspore/ccsrc/transform/convert.cc b/mindspore/ccsrc/transform/convert.cc index cfbcdb0bfd0..a8c085c5fb9 100644 --- a/mindspore/ccsrc/transform/convert.cc +++ b/mindspore/ccsrc/transform/convert.cc @@ -205,6 +205,8 @@ const char kNameL2Loss[] = "L2Loss"; const char kNameCTCLoss[] = "CTCLoss"; const char kNameRange[] = "Range"; const char kNameSquareSumAll[] = "SquareSumAll"; +const char kNameAscendQuant[] = "AscendQuant"; +const char kNameAscendDequant[] = "AscendDequant"; // -----------------OpAdapter initialization-------------- std::unordered_map &DfGraphConvertor::get_adpt_map() { @@ -408,7 +410,9 @@ std::unordered_map &DfGraphConvertor::get_adpt_ma {string(kNameL2Loss), ADPT_DESC(L2Loss)}, {string(kNameCTCLoss), ADPT_DESC(CTCLoss)}, {string(kNameRange), ADPT_DESC(RangeD)}, - {string(kNameSquareSumAll), ADPT_DESC(SquareSumAll)}}; + {string(kNameSquareSumAll), ADPT_DESC(SquareSumAll)}, + {string(kNameAscendQuant), ADPT_DESC(AscendQuant)}, + {string(kNameAscendDequant), ADPT_DESC(AscendDequant)}}; #ifdef ENABLE_GE adpt_map[string(kNamePrint)] = ADPT_DESC(Print); adpt_map[string(kNameApplyAdam)] = ADPT_DESC(ApplyAdamD); diff --git a/mindspore/ccsrc/transform/op_declare.cc b/mindspore/ccsrc/transform/op_declare.cc index 7643fa54e4f..64ec0626104 100644 --- a/mindspore/ccsrc/transform/op_declare.cc +++ b/mindspore/ccsrc/transform/op_declare.cc @@ -1277,6 +1277,19 @@ ATTR_MAP(CTCLoss) = { {"ignore_longer_outputs_than_inputs", ATTR_DESC(ignore_longer_outputs_than_inputs, AnyTraits())}}; OUTPUT_MAP(CTCLoss) = {{0, OUTPUT_DESC(loss)}, {1, OUTPUT_DESC(gradient)}}; +// AscendQuant +INPUT_MAP(AscendQuant) = {{1, INPUT_DESC(x)}}; +ATTR_MAP(AscendQuant) = {{"scale", ATTR_DESC(scale, AnyTraits())}, + {"offset", ATTR_DESC(offset, AnyTraits())}, + {"sqrt_mode", ATTR_DESC(sqrt_mode, AnyTraits())}, + {"round_mode", ATTR_DESC(round_mode, AnyTraits())}}; +OUTPUT_MAP(AscendQuant) = {{0, OUTPUT_DESC(y)}}; + +// AscendDequant +INPUT_MAP(AscendDequant) = {{1, INPUT_DESC(x)}, {2, INPUT_DESC(deq_scale)}}; +ATTR_MAP(AscendDequant) = {{"sqrt_mode", ATTR_DESC(sqrt_mode, AnyTraits())}, + {"relu_flag", ATTR_DESC(relu_flag, AnyTraits())}}; +OUTPUT_MAP(AscendDequant) = {{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 6e0debd572d..6edc0e1884f 100755 --- a/mindspore/ccsrc/transform/op_declare.h +++ b/mindspore/ccsrc/transform/op_declare.h @@ -481,6 +481,10 @@ DECLARE_OP_ADAPTER(L2Loss) DECLARE_OP_USE_OUTPUT(L2Loss) DECLARE_OP_ADAPTER(CTCLoss) DECLARE_OP_USE_OUTPUT(CTCLoss) +DECLARE_OP_ADAPTER(AscendQuant) +DECLARE_OP_USE_OUTPUT(AscendQuant) +DECLARE_OP_ADAPTER(AscendDequant) +DECLARE_OP_USE_OUTPUT(AscendDequant) #ifdef ENABLE_GE DECLARE_OP_ADAPTER(Print) DECLARE_OP_USE_DYN_INPUT(Print) diff --git a/mindspore/ops/operations/_quant_ops.py b/mindspore/ops/operations/_quant_ops.py index b228c51b10a..5fb92dee659 100644 --- a/mindspore/ops/operations/_quant_ops.py +++ b/mindspore/ops/operations/_quant_ops.py @@ -39,6 +39,8 @@ __all__ = ["FakeQuantPerLayer", "BatchNormFold2_D", "BatchNormFold2GradD", "BatchNormFold2GradReduce", + "AscendQuant", + "AscendDequant", ] @@ -975,3 +977,104 @@ class FakeQuantMinMaxPerChannelUpdate(PrimitiveWithInfer): validator.check_tensor_type_same( {"max": max_type}, valid_types, self.name) return min_type, max_type + + +class AscendQuant(PrimitiveWithInfer): + r""" + Returns the quantized value of input_x. + + If `sqrt_mode` is False: + + .. math:: + y = round(scale * x + offset) + If `sqrt_mode` is True: + + .. math:: + y = round(scale * x * scale + offset) + + Note: + This operation only support Ascend 310 inference environment. + + Args: + scale (float) : Specifies the scaling ratio. + offset (float): Specifies the offset. + sqrt_mode (bool) : Specifies whether to perform square root on `scale`. Default: False. + round_mode (str): Specifies the way to round. Should be one of ["Round", "Floor", "Ceil", "Trunc"]. + Default: "Round". + + Inputs: + - **input_x** (Tensor) : Input tensor. Its data type should be mindspore.float16 or mindspore.float32. + + Outputs: + - Tensor: The quantized output tensor of type mindspore.int8. + + Examples: + >>> input_x = Tensor([100.0, 150.0], mstype.float32) + >>> quant = P.AscendQuant(80.0, 0.0, False, "Round") + >>> y = quant(input_x) + """ + + @prim_attr_register + def __init__(self, scale, offset, sqrt_mode=False, round_mode="Round"): + self.scale = validator.check_value_type("scale", scale, [float], self.name) + self.offset = validator.check_value_type("offset", offset, [float], self.name) + self.sqrt_mode = validator.check_value_type("sqrt_mode", sqrt_mode, [bool], self.name) + self.round_mode = validator.check_string("round_mode", round_mode, + ["Round", "Floor", "Ceil", "Trunc"], self.name) + + def infer_shape(self, x_shape): + return x_shape + + def infer_dtype(self, x_type): + validator.check_subclass("input_x", x_type, mstype.tensor, self.name) + validator.check_type_name("input_x", x_type, [mstype.float16, mstype.float32], self.name) + return mstype.int8 + + +class AscendDequant(PrimitiveWithInfer): + r""" + Returns the dequantized value of input_x. + This operation will do ReLU to the dequantized value if `relu_flag` is True. + + If `sqrt_mode` is False: + + .. math:: + y = x * deq\_scale + If `sqrt_mode` is True: + + .. math:: + y = x * deq\_scale * deq\_scale + + Note: + This operation only support Ascend 310 inference environment. + + Args: + sqrt_mode (bool) : Specifies whether to perform square root on `scale`. Default: False. + relu_flag (bool): Specifies whether to perform ReLU. Default: False. + + Inputs: + - **input_x** (Tensor) : Input tensor. Should be mindspore.int32. + - **deq_scale** (Tensor) : Specifies the scaling ratio. + Data type should be mindspore.float16 or mindspore.uint64 + + Outputs: + - Tensor: The quantized output tensor of type mindspore.float16. + + Examples: + >>> input_x = Tensor([100.0, 150.0], mstype.float32) + >>> dequant = P.AscendDequant(False, False) + >>> y = dequant(input_x) + """ + @prim_attr_register + def __init__(self, sqrt_mode=False, relu_flag=False): + self.sqrt_mode = validator.check_value_type("sqrt_mode", sqrt_mode, [bool], self.name) + self.relu_flag = validator.check_value_type("relu_flag", relu_flag, [bool], self.name) + + def infer_shape(self, x_shape, deq_scale_shape): + return x_shape + + def infer_dtype(self, x_type, deq_scale_type): + validator.check_subclass("x", x_type, mstype.tensor, self.name) + validator.check_type_name("x", x_type, [mstype.int32], self.name) + validator.check_type_name("deq_scale", deq_scale_type, [mstype.float16, mstype.uint64], self.name) + return mstype.float16 diff --git a/tests/ut/python/ops/test_ops.py b/tests/ut/python/ops/test_ops.py index 03d24375d55..42f283448ff 100755 --- a/tests/ut/python/ops/test_ops.py +++ b/tests/ut/python/ops/test_ops.py @@ -1614,7 +1614,43 @@ test_case_other_ops = [ ] -test_case_lists = [test_case_nn_ops, test_case_math_ops, test_case_array_ops, test_case_other_ops] + +test_case_quant_ops = [ + ('AscendQuant_1', { + 'block': P.AscendQuant(0.5, 0.0, False, "Round"), + 'desc_inputs': [Tensor(np.random.rand(1,2,4,4), mstype.float32)], + 'skip': ['backward']}), + ('AscendQuant_2', { + 'block': P.AscendQuant(80.0, 10.0, True, "Round"), + 'desc_inputs': [Tensor([100.0, 200.0], mstype.float32)], + 'skip': ['backward']}), + ('AscendQuant_3', { + 'block': P.AscendQuant(80.0, 0.0, False, "Floor"), + 'desc_inputs': [Tensor([100.0, 200.0], mstype.float32)], + 'skip': ['backward']}), + ('AscendQuant_4', { + 'block': P.AscendQuant(80.0, 0.0, False, "Ceil"), + 'desc_inputs': [Tensor([100.0, 200.0], mstype.float32)], + 'skip': ['backward']}), + ('AscendQuant_5', { + 'block': P.AscendQuant(80.0, 0.0, False, "Trunc"), + 'desc_inputs': [Tensor([100.0, 200.0], mstype.float32)], + 'skip': ['backward']}), + ('AscendQuant_6', { + 'block': P.AscendQuant(-80.0, 10.0, False, "Round"), + 'desc_inputs': [Tensor([100.0, 200.0], mstype.float32)], + 'skip': ['backward']}), + ('AscendQuant_7', { + 'block': P.AscendQuant(80.0, -10.0, False, "Round"), + 'desc_inputs': [Tensor([100.0, 200.0], mstype.float32)], + 'skip': ['backward']}), + ('AscendQuant_8', { + 'block': P.AscendQuant(80.0, 10.0, False, "Round"), + 'desc_inputs': [Tensor([100.0, 200.0], mstype.float16)], + 'skip': ['backward']}), +] + +test_case_lists = [test_case_nn_ops, test_case_math_ops, test_case_array_ops, test_case_other_ops, test_case_quant_ops] test_case = functools.reduce(lambda x, y: x + y, test_case_lists) # use -k to select certain testcast # pytest tests/python/ops/test_ops.py::test_backward -k LayerNorm