diff --git a/docs/OperatorKernels.md b/docs/OperatorKernels.md index ee92269f1b..c3e2ff94d3 100644 --- a/docs/OperatorKernels.md +++ b/docs/OperatorKernels.md @@ -25,7 +25,7 @@ Do not modify directly.* |||[7, 12]|**T** = tensor(double), tensor(float), tensor(int32), tensor(int64)| |Affine|*in* X:**T**
*out* Y:**T**|1+|**T** = tensor(float)| |And|*in* A:**T**
*in* B:**T**
*out* C:**T1**|7+|**T** = tensor(bool)
**T1** = tensor(bool)| -|ArgMax|*in* data:**T**
*out* reduced:**tensor(int64)**|13+|**T** = tensor(double), tensor(float), tensor(int32)| +|ArgMax|*in* data:**T**
*out* reduced:**tensor(int64)**|13+|**T** = tensor(double), tensor(float), tensor(int32), tensor(int8), tensor(uint8)| |||[11, 12]|**T** = tensor(double), tensor(float), tensor(int32)| |||[1, 10]|**T** = tensor(float), tensor(int32)| |ArgMin|*in* data:**T**
*out* reduced:**tensor(int64)**|13+|**T** = tensor(double), tensor(float), tensor(int32)| diff --git a/onnxruntime/core/providers/cpu/cpu_execution_provider.cc b/onnxruntime/core/providers/cpu/cpu_execution_provider.cc index f015b53941..eb12b7c4fd 100644 --- a/onnxruntime/core/providers/cpu/cpu_execution_provider.cc +++ b/onnxruntime/core/providers/cpu/cpu_execution_provider.cc @@ -526,6 +526,8 @@ class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 13, double_double, Dropout); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 13, float, ArgMax); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 13, double, ArgMax); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 13, int8_t, ArgMax); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 13, uint8_t, ArgMax); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 13, int32_t, ArgMax); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 13, float, ArgMin); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 13, double, ArgMin); @@ -1570,6 +1572,10 @@ Status RegisterOnnxOperatorKernels(KernelRegistry& kernel_registry) { float, ArgMax)>, BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo{0}); test.AddAttribute("keepdims", (int64_t)0); test.AddInput("data", {1, 2, 2}, - {1.0f, 2.0f, - 3.0f, 4.0f}); + {1.0f, 2.0f, + 3.0f, 4.0f}); test.AddOutput("reduced", {2, 2}, {1.f, 2.f, 3.f, 4.f}); test.Run(); } @@ -2181,6 +2181,48 @@ TEST(ReductionOpTest, ArgMax_int32_neg_axis) { test.Run(); } +TEST(ReductionOpTest, ArgMax_int8) { + OpTester test("ArgMax", 13); + test.AddAttribute("axis", static_cast(1)); + test.AddAttribute("keepdims", static_cast(1)); + test.AddInput("data", {3, 2, 2}, + {1, 2, + 3, 4, + + 5, 6, + 7, 8, + + 9, 10, + 11, 12}); + test.AddOutput("reduced", {3, 1, 2}, + {1, 1, + 1, 1, + 1, 1}); + // TensorRT: input/output with DataType Int8 in network without Q/DQ layers + // must have dynamic range set when no calibrator is used + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kTensorrtExecutionProvider}); +} + +TEST(ReductionOpTest, ArgMax_uint8) { + OpTester test("ArgMax", 13); + test.AddAttribute("axis", static_cast(1)); + test.AddAttribute("keepdims", static_cast(1)); + test.AddInput("data", {3, 2, 2}, + {1, 2, + 3, 4, + + 5, 6, + 7, 8, + + 9, 10, + 11, 12}); + test.AddOutput("reduced", {3, 1, 2}, + {1, 1, + 1, 1, + 1, 1}); + test.Run(); +} + TEST(ReductionOpTest, ArgMax2D) { OpTester test("ArgMax"); test.AddAttribute("axis", (int64_t)1); diff --git a/onnxruntime/test/python/quantization/test_op_argmax.py b/onnxruntime/test/python/quantization/test_op_argmax.py new file mode 100644 index 0000000000..cb0a243c7e --- /dev/null +++ b/onnxruntime/test/python/quantization/test_op_argmax.py @@ -0,0 +1,116 @@ +#!/usr/bin/env python +# coding: utf-8 +# ------------------------------------------------------------------------- +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. See License.txt in the project root for +# license information. +# -------------------------------------------------------------------------- + +import unittest +import onnx +import numpy as np +from onnx import helper, TensorProto +from onnxruntime.quantization import quantize_static, QuantFormat, QuantType +from op_test_utils import TestDataFeeds, check_model_correctness, check_op_type_count, check_op_nodes, check_qtype_by_node_type + + +class TestOpArgMax(unittest.TestCase): + def input_feeds(self, n, name2shape): + input_data_list = [] + for i in range(n): + inputs = {} + for name, shape in name2shape.items(): + inputs.update({name: np.random.randint(-1, 2, shape).astype(np.float32)}) + input_data_list.extend([inputs]) + dr = TestDataFeeds(input_data_list) + return dr + + def construct_model_argmax(self, output_model_path, input_shape, output_shape): + # (input) + # | + # Conv + # | + # ArgMax + # | + # (output) + input_name = 'input' + output_name = 'output' + initializers = [] + + # make Conv node + conv_weight_name = 'conv_weight' + conv_weight_arr = np.random.randint(-1, 2, [32, 256, 1, 1]).astype(np.float32) + conv_weight_initializer = onnx.numpy_helper.from_array(conv_weight_arr, name=conv_weight_name) + conv_output_name = 'conv_output' + conv_inputs = [input_name, conv_weight_name] + conv_outputs = [conv_output_name] + conv_name = 'conv_node' + conv_node = onnx.helper.make_node('Conv', conv_inputs, conv_outputs, dilations=[1, 1], kernel_shape=[1, 1], + pads=[0, 0, 0, 0], strides=[1, 1], name=conv_name) + + # make ArgMax node + argmax_inputs = [conv_output_name] + argmax_outputs = [output_name] + argmax_name = 'argmax_node' + argmax_node = onnx.helper.make_node('ArgMax', argmax_inputs, argmax_outputs, axis=3, keepdims=0, name=argmax_name) + + initializers = [conv_weight_initializer] + + # make graph + input_tensor = helper.make_tensor_value_info(input_name, TensorProto.FLOAT, input_shape) + output_tensor = helper.make_tensor_value_info(output_name, TensorProto.INT64, output_shape) + graph_name = 'ArgMax_Quant_Test' + graph = helper.make_graph([conv_node, argmax_node], graph_name, + [input_tensor], [output_tensor], initializer=initializers) + model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 13)]) + model.ir_version = 7 # use stable onnx ir version + + onnx.save(model, output_model_path) + + def quantize_argmax_test(self, activation_type, weight_type, extra_options = {}): + np.random.seed(1) + model_fp32_path = 'argmax_fp32.onnx' + + self.construct_model_argmax(model_fp32_path, + [1, 256, 128, 128], + [1, 32, 128]) + + activation_proto_qtype = TensorProto.UINT8 if activation_type == QuantType.QUInt8 else TensorProto.INT8 + activation_type_str = 'u8' if (activation_type == QuantType.QUInt8) else 's8' + weight_type_str = 'u8' if (weight_type == QuantType.QUInt8) else 's8' + model_uint8_path = 'argmax_{}{}.onnx'.format(activation_type_str, weight_type_str) + model_uint8_qdq_path = 'argmax_{}{}_qdq.onnx'.format(activation_type_str, weight_type_str) + + # Verify QOperator mode + data_reader = self.input_feeds(1, {'input': [1, 256, 128, 128]}) + quantize_static(model_fp32_path, model_uint8_path, data_reader, + activation_type = activation_type, weight_type = weight_type, extra_options = extra_options) + # make sure argmax become xint8 operator, its input name could tell that + check_op_nodes(self, model_uint8_path, lambda node: not(node.name == "argmax_node" and node.input[0] == 'conv_output')) + qnode_counts = {'QuantizeLinear': 1, 'QLinearConv': 1, 'ArgMax': 1} + check_op_type_count(self, model_uint8_path, **qnode_counts) + qnode_io_qtypes = {'QuantizeLinear' : [['i', 2, activation_proto_qtype], ['o', 0, activation_proto_qtype]]} + check_qtype_by_node_type(self, model_uint8_path, qnode_io_qtypes) + data_reader.rewind() + check_model_correctness(self, model_fp32_path, model_uint8_path, data_reader.get_next()) + + # Verify QDQ mode + data_reader.rewind() + quantize_static(model_fp32_path, model_uint8_qdq_path, data_reader, quant_format=QuantFormat.QDQ, + activation_type = activation_type, weight_type = weight_type, extra_options = extra_options) + qdqnode_counts = {'QuantizeLinear': 2, 'DequantizeLinear': 3, 'ArgMax': 1} + check_op_type_count(self, model_uint8_qdq_path, **qdqnode_counts) + qnode_io_qtypes = {'QuantizeLinear' : [['i', 2, activation_proto_qtype], ['o', 0, activation_proto_qtype]]} + check_qtype_by_node_type(self, model_uint8_qdq_path, qnode_io_qtypes) + data_reader.rewind() + check_model_correctness(self, model_fp32_path, model_uint8_qdq_path, data_reader.get_next()) + + + def test_quantize_argmax(self): + self.quantize_argmax_test(QuantType.QUInt8, QuantType.QUInt8) + + def test_quantize_argmax_s8s8(self): + self.quantize_argmax_test(QuantType.QInt8, QuantType.QInt8, extra_options = {'ActivationSymmetric' : True}) + +if __name__ == '__main__': + unittest.main() diff --git a/onnxruntime/test/testdata/kernel_def_hashes/onnx.cpu.json b/onnxruntime/test/testdata/kernel_def_hashes/onnx.cpu.json index 41027eff75..88974c2b26 100644 --- a/onnxruntime/test/testdata/kernel_def_hashes/onnx.cpu.json +++ b/onnxruntime/test/testdata/kernel_def_hashes/onnx.cpu.json @@ -139,6 +139,10 @@ "And ai.onnx CPUExecutionProvider", 7931711152704979424 ], + [ + "ArgMax ai.onnx CPUExecutionProvider", + 1506815612154586624 + ], [ "ArgMax ai.onnx CPUExecutionProvider", 2317240841713684752 @@ -163,6 +167,10 @@ "ArgMax ai.onnx CPUExecutionProvider", 12157492629288967928 ], + [ + "ArgMax ai.onnx CPUExecutionProvider", + 12195502223979275536 + ], [ "ArgMax ai.onnx CPUExecutionProvider", 13027735142126919896