saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-29 20:14:01 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Quantization of Argmax (#10213)

Browse source 
This patch includes:
* int8/uint8 support for Argmax
* Quantization tool support for Argmax
This commit is contained in:
Yi-Hong Lyu 2022-01-13 06:12:56 +08:00 committed by GitHub
parent 98f85ae05b
commit 499f1d5fd7
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
8 changed files with 196 additions and 3 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
docs/OperatorKernels.md
View file

@ -25,7 +25,7 @@ Do not modify directly.*
|||[7, 12]|**T** = tensor(double), tensor(float), tensor(int32), tensor(int64)|
|Affine|*in* X:**T**<br> *out* Y:**T**|1+|**T** = tensor(float)|
|And|*in* A:**T**<br> *in* B:**T**<br> *out* C:**T1**|7+|**T** = tensor(bool)<br/> **T1** = tensor(bool)|
|ArgMax|*in* data:**T**<br> *out* reduced:**tensor(int64)**|13+|**T** = tensor(double), tensor(float), tensor(int32)|
|ArgMax|*in* data:**T**<br> *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**<br> *out* reduced:**tensor(int64)**|13+|**T** = tensor(double), tensor(float), tensor(int32)|

6
onnxruntime/core/providers/cpu/cpu_execution_provider.cc
View file

@ -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<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 13,
double, ArgMax)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 13,
int8_t, ArgMax)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 13,
uint8_t, ArgMax)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 13,
int32_t, ArgMax)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 13,

2
onnxruntime/core/providers/cpu/reduction/reduction_ops.cc
View file

@ -203,6 +203,8 @@ REGISTER_UNARY_ELEMENTWISE_VERSIONED_KERNEL(ArgMax, 11, 12);
REGISTER_UNARY_ELEMENTWISE_VERSIONED_KERNEL_DOUBLE_ONLY(ArgMax, 11, 12)
REGISTER_UNARY_ELEMENTWISE_KERNEL(ArgMax, 13);
REGISTER_UNARY_ELEMENTWISE_KERNEL_DOUBLE_ONLY(ArgMax, 13);
REGISTER_UNARY_ELEMENTWISE_KERNEL_INT8_ONLY(ArgMax, 13);
REGISTER_UNARY_ELEMENTWISE_KERNEL_UINT8_ONLY(ArgMax, 13);
REGISTER_UNARY_ELEMENTWISE_VERSIONED_KERNEL(ArgMin, 1, 10);
REGISTER_UNARY_ELEMENTWISE_VERSIONED_KERNEL(ArgMin, 11, 12);

17
onnxruntime/python/tools/quantization/operators/argmax.py Normal file
View file

@ -0,0 +1,17 @@
from .base_operator import QuantOperatorBase
# Use the quantized tensor as input without DQ.
class QArgMax(QuantOperatorBase):
def __init__(self, onnx_quantizer, onnx_node):
super().__init__(onnx_quantizer, onnx_node)
def quantize(self):
node = self.node
quantized_input_value = self.quantizer.find_quantized_value(node.input[0])
if quantized_input_value is None:
self.quantizer.new_nodes += [node]
return
node.input[0] = quantized_input_value.q_name
self.quantizer.new_nodes += [node]

2
onnxruntime/python/tools/quantization/registry.py
View file

@ -1,4 +1,5 @@
from .quant_utils import QuantizationMode
from .operators.argmax import QArgMax
from .operators.base_operator import QuantOperatorBase
from .operators.qdq_base_operator import QDQOperatorBase
from .operators.matmul import MatMulInteger, QLinearMatMul, QDQMatMul
@ -33,6 +34,7 @@ IntegerOpsRegistry = {
IntegerOpsRegistry.update(CommonOpsRegistry)
QLinearOpsRegistry = {
"ArgMax": QArgMax,
"Conv": QLinearConv,
"MatMul": QLinearMatMul,
"Add": QLinearBinaryOp,

46
onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc
View file

@ -562,8 +562,8 @@ TEST(ReductionOpTest, ReduceLogSumExp_float_no_reduction) {
test.AddAttribute("axes", std::vector<int64_t>{0});
test.AddAttribute("keepdims", (int64_t)0);
test.AddInput<float>("data", {1, 2, 2},
{1.0f, 2.0f,
3.0f, 4.0f});
{1.0f, 2.0f,
3.0f, 4.0f});
test.AddOutput<float>("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<int64_t>(1));
test.AddAttribute("keepdims", static_cast<int64_t>(1));
test.AddInput<int8_t>("data", {3, 2, 2},
{1, 2,
3, 4,
5, 6,
7, 8,
9, 10,
11, 12});
test.AddOutput<int64_t>("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<int64_t>(1));
test.AddAttribute("keepdims", static_cast<int64_t>(1));
test.AddInput<uint8_t>("data", {3, 2, 2},
{1, 2,
3, 4,
5, 6,
7, 8,
9, 10,
11, 12});
test.AddOutput<int64_t>("reduced", {3, 1, 2},
{1, 1,
1, 1,
1, 1});
test.Run();
}
TEST(ReductionOpTest, ArgMax2D) {
OpTester test("ArgMax");
test.AddAttribute("axis", (int64_t)1);

116
onnxruntime/test/python/quantization/test_op_argmax.py Normal file
View file

@ -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()

8
onnxruntime/test/testdata/kernel_def_hashes/onnx.cpu.json vendored
View file

@ -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