mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-29 20:14:01 +00:00
[QNN EP] Add support for GatherElements (#15966)
### Description - Adds support for the GatherElements operator to QNN EP. - Adds GatherElements to QDQ quantizer tool. ### Motivation and Context Enable more models to run on QNN EP.
This commit is contained in:
Adrian Lizarraga
GitHub
parent
7c93d5ded1
commit
a22cc078b4
11 changed files with 609 additions and 62 deletions
|
|
@ -33,6 +33,7 @@ void Selectors::RegisterSelector(const OpVersionsAndSelector::OpVersionsMap& ops
|
|||
// output Q have the same scale and zero_point.
|
||||
static const OpVersionsAndSelector::OpVersionsMap GetMiscOpVersionsMap() {
|
||||
return {{"Gather", {}},
|
||||
{"GatherElements", {}},
|
||||
{"Reshape", {}},
|
||||
{"Expand", {}},
|
||||
{"Flatten", {}},
|
||||
|
|
|
|||
|
|
@ -110,6 +110,7 @@ OpBuilderRegistrations::OpBuilderRegistrations() {
|
|||
|
||||
{
|
||||
CreateGatherOpBuilder("Gather", *this);
|
||||
CreateGatherOpBuilder("GatherElements", *this);
|
||||
}
|
||||
|
||||
{
|
||||
|
|
|
|||
|
|
@ -122,6 +122,7 @@ class BaseOpBuilder : public IOpBuilder {
|
|||
{"Exp", QNN_OP_ELEMENT_WISE_EXP},
|
||||
{"Floor", QNN_OP_ELEMENT_WISE_FLOOR},
|
||||
{"Gather", QNN_OP_GATHER},
|
||||
{"GatherElements", QNN_OP_GATHER_ELEMENTS},
|
||||
{"Greater", QNN_OP_ELEMENT_WISE_GREATER},
|
||||
{"GreaterOrEqual", QNN_OP_ELEMENT_WISE_GREATER_EQUAL},
|
||||
{"Less", QNN_OP_ELEMENT_WISE_LESS},
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
// Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
// Licensed under the MIT License.
|
||||
|
||||
#include <cassert>
|
||||
#include "core/providers/common.h"
|
||||
#include "core/providers/shared/utils/utils.h"
|
||||
#include "core/providers/qnn/builder/qnn_model_wrapper.h"
|
||||
|
|
@ -13,11 +14,16 @@
|
|||
namespace onnxruntime {
|
||||
namespace qnn {
|
||||
|
||||
// Handles Gather and GatherElements
|
||||
class GatherOpBuilder : public BaseOpBuilder {
|
||||
public:
|
||||
GatherOpBuilder() : BaseOpBuilder("GatherOpBuilder") {}
|
||||
ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(GatherOpBuilder);
|
||||
|
||||
Status IsOpSupported(QnnModelWrapper& qnn_model_wrapper,
|
||||
const NodeUnit& node_unit,
|
||||
const logging::Logger& logger) const override ORT_MUST_USE_RESULT;
|
||||
|
||||
protected:
|
||||
Status ProcessInputs(QnnModelWrapper& qnn_model_wrapper,
|
||||
const NodeUnit& node_unit,
|
||||
|
|
@ -32,79 +38,154 @@ class GatherOpBuilder : public BaseOpBuilder {
|
|||
bool do_op_validation) const override ORT_MUST_USE_RESULT;
|
||||
};
|
||||
|
||||
Status GatherOpBuilder::ProcessInputs(QnnModelWrapper& qnn_model_wrapper,
|
||||
Status GatherOpBuilder::IsOpSupported(QnnModelWrapper& qnn_model_wrapper,
|
||||
const NodeUnit& node_unit,
|
||||
const logging::Logger& logger,
|
||||
std::vector<std::string>& input_names,
|
||||
bool do_op_validation) const {
|
||||
const auto& inputs = node_unit.Inputs();
|
||||
ORT_RETURN_IF(inputs.size() != 2, "Gather should has 2 inputs at least!");
|
||||
ORT_RETURN_IF_ERROR(ProcessInput(qnn_model_wrapper, inputs[0], logger, input_names));
|
||||
const logging::Logger& logger) const {
|
||||
// On QNN CPU backend, the QNN validator does not properly reject unsupported input shapes.
|
||||
// This causes a Qnn graph execution error. So, reject those configs here.
|
||||
// We should consider not using QNN CPU backend for onnxruntime unit tests.
|
||||
const std::string& op_type = node_unit.OpType();
|
||||
if (qnn_model_wrapper.GetQnnBackendType() == QnnBackendType::CPU && op_type == "GatherElements") {
|
||||
const auto& input0 = node_unit.Inputs()[0];
|
||||
std::vector<uint32_t> input0_shape;
|
||||
ORT_RETURN_IF_NOT(qnn_model_wrapper.GetOnnxShape(input0.node_arg, input0_shape),
|
||||
"Cannot get input[0] shape for ", op_type, " node ", node_unit.Name());
|
||||
|
||||
// Process indices
|
||||
const auto& input_name = inputs[1].node_arg.Name();
|
||||
const size_t input0_rank = input0_shape.size();
|
||||
ORT_RETURN_IF_NOT(input0_rank > 1 && input0_rank <= 4,
|
||||
"QNN CPU backend does not support ", op_type, " with input[0] of rank ", input0_rank);
|
||||
}
|
||||
|
||||
return BaseOpBuilder::IsOpSupported(qnn_model_wrapper, node_unit, logger);
|
||||
}
|
||||
|
||||
// Makes negative indices positive and converts int64 indices to another integer type (typically int32 or uint32).
|
||||
// The input and output are both represented as byte arrays.
|
||||
template <typename SrcType, typename DstType>
|
||||
static bool FixStaticIndices(const std::vector<uint8_t>& onnx_bytes,
|
||||
int64_t input0_axis_dim,
|
||||
/*out*/ std::vector<uint8_t>& qnn_bytes) {
|
||||
const size_t num_elems = onnx_bytes.size() / sizeof(SrcType);
|
||||
gsl::span<const SrcType> onnx_indices{reinterpret_cast<const SrcType*>(onnx_bytes.data()), num_elems};
|
||||
|
||||
qnn_bytes.resize(num_elems * sizeof(DstType));
|
||||
gsl::span<DstType> qnn_indices{reinterpret_cast<DstType*>(qnn_bytes.data()), num_elems};
|
||||
|
||||
for (size_t i = 0; i < num_elems; i++) {
|
||||
SrcType onnx_index = onnx_indices[i];
|
||||
|
||||
// Try to make a negative index positive by adding rank.
|
||||
if (onnx_index < 0) {
|
||||
onnx_index += static_cast<SrcType>(input0_axis_dim);
|
||||
}
|
||||
|
||||
if (onnx_index < 0 || static_cast<int64_t>(onnx_index) >= input0_axis_dim) {
|
||||
return false; // QNN does not support out-of-bounds indices.
|
||||
}
|
||||
|
||||
qnn_indices[i] = static_cast<DstType>(onnx_index);
|
||||
}
|
||||
|
||||
return true;
|
||||
}
|
||||
|
||||
// Gets the size of input0 on the axis dimension.
|
||||
static Status GetInpu0AxisDimValue(const QnnModelWrapper& qnn_model_wrapper,
|
||||
const NodeUnit& node_unit,
|
||||
int64_t default_axis_value,
|
||||
/*out*/ int64_t& axis_dim_value) {
|
||||
const auto& input0 = node_unit.Inputs()[0];
|
||||
std::vector<uint32_t> input0_shape;
|
||||
ORT_RETURN_IF_NOT(qnn_model_wrapper.GetOnnxShape(input0.node_arg, input0_shape),
|
||||
"Cannot get shape for ", node_unit.OpType(), " input[0] ", input0.node_arg.Name());
|
||||
|
||||
int64_t rank = static_cast<int64_t>(input0_shape.size());
|
||||
NodeAttrHelper node_helper(node_unit);
|
||||
int64_t onnx_axis = node_helper.Get("axis", default_axis_value);
|
||||
if (onnx_axis < 0) {
|
||||
onnx_axis += rank;
|
||||
}
|
||||
ORT_RETURN_IF_NOT((onnx_axis >= 0 && onnx_axis < static_cast<int64_t>(input0_shape.size())),
|
||||
"QNN requires axis range [0, rank-1] for ", node_unit.OpType());
|
||||
|
||||
axis_dim_value = static_cast<int64_t>(input0_shape[onnx_axis]);
|
||||
|
||||
return Status::OK();
|
||||
}
|
||||
|
||||
// Processes the indices input to Gather operators.
|
||||
//
|
||||
// In general, QNN only supports int32/uint32 indices. QNN EP has to add Cast for dynamic int64 indices or
|
||||
// convert static int64 indices to int32.
|
||||
//
|
||||
// The HTP backend only supports dynamic int64 indices if they are a graph input.
|
||||
static Status ProcessIndicesInput(QnnModelWrapper& qnn_model_wrapper,
|
||||
const NodeUnitIODef& indices_input,
|
||||
int64_t input0_axis_dim,
|
||||
const logging::Logger& logger,
|
||||
std::vector<std::string>& input_names,
|
||||
bool do_op_validation) {
|
||||
const auto& input_name = indices_input.node_arg.Name();
|
||||
if (qnn_model_wrapper.IsQnnTensorWrapperExist(input_name)) {
|
||||
LOGS(logger, VERBOSE) << "Tensor already added, skip it: " << input_name;
|
||||
input_names.push_back(input_name);
|
||||
return Status::OK();
|
||||
}
|
||||
|
||||
std::string indices_input_name(input_name);
|
||||
Qnn_DataType_t qnn_data_type = QNN_DATATYPE_INT_32;
|
||||
const auto* type_proto = inputs[1].node_arg.TypeAsProto();
|
||||
ORT_RETURN_IF_ERROR(utils::GetQnnDataType(false, type_proto, qnn_data_type));
|
||||
TensorInfo indices_info = {};
|
||||
ORT_RETURN_IF_ERROR(qnn_model_wrapper.GetTensorInfo(indices_input, indices_info));
|
||||
|
||||
std::vector<uint8_t> unpacked_tensor;
|
||||
std::vector<uint8_t> gather_indices;
|
||||
bool is_initializer_input = qnn_model_wrapper.IsInitializerInput(input_name);
|
||||
const bool is_npu_backend = IsNpuBackend(qnn_model_wrapper.GetQnnBackendType());
|
||||
const bool is_graph_input = qnn_model_wrapper.IsGraphInput(input_name);
|
||||
ORT_RETURN_IF(is_npu_backend &&
|
||||
(indices_info.qnn_data_type == QNN_DATATYPE_INT_64) &&
|
||||
!(indices_info.is_initializer || is_graph_input),
|
||||
"HTP backend doesn't support a Gather* op with a dynamic int64 input activation ",
|
||||
"unless it is a graph input.");
|
||||
|
||||
// Gather input 0 is quantized tensor, input 1 (indices) is int64, this is not supported by QNN
|
||||
bool is_quantized_tensor = inputs[0].quant_param.has_value();
|
||||
ORT_RETURN_IF(is_quantized_tensor && qnn_data_type == QNN_DATATYPE_INT_64 && !is_initializer_input,
|
||||
"HTP backend doesn't support any int64 data type.");
|
||||
std::vector<uint8_t> qnn_indices_bytes;
|
||||
|
||||
if (is_initializer_input) {
|
||||
const auto& input_tensor = qnn_model_wrapper.GetInitializerTensors().at(input_name);
|
||||
ORT_RETURN_IF_ERROR(qnn_model_wrapper.UnpackInitializerData(*input_tensor, unpacked_tensor));
|
||||
if (qnn_data_type == QNN_DATATYPE_INT_64) {
|
||||
// Convert initializer from int64 to int32
|
||||
size_t size = unpacked_tensor.size() / sizeof(int64_t);
|
||||
const int64_t* gather_indices_int64 = reinterpret_cast<const int64_t*>(unpacked_tensor.data());
|
||||
gather_indices.resize(size * sizeof(int32_t));
|
||||
int32_t* gather_indices_int32 = reinterpret_cast<int32_t*>(gather_indices.data());
|
||||
std::transform(gather_indices_int64, gather_indices_int64 + size, gather_indices_int32,
|
||||
[](int64_t item) { return SafeInt<uint32_t>(item); });
|
||||
// Get raw bytes for static indices.
|
||||
// If indices are int64, convert them to int32 and update indices_info.qnn_data_type.
|
||||
if (indices_info.is_initializer) {
|
||||
std::vector<uint8_t> onnx_indices_bytes;
|
||||
ORT_RETURN_IF_ERROR(qnn_model_wrapper.UnpackInitializerData(*indices_info.initializer_tensor, onnx_indices_bytes));
|
||||
|
||||
if (indices_info.qnn_data_type == QNN_DATATYPE_INT_64) {
|
||||
ORT_RETURN_IF_NOT((FixStaticIndices<int64_t, int32_t>(onnx_indices_bytes, input0_axis_dim, qnn_indices_bytes)),
|
||||
"QNN does not support negative index values for Gather* ops");
|
||||
indices_info.qnn_data_type = QNN_DATATYPE_INT_32;
|
||||
} else if (indices_info.qnn_data_type == QNN_DATATYPE_INT_32) {
|
||||
ORT_RETURN_IF_NOT((FixStaticIndices<int32_t, int32_t>(onnx_indices_bytes, input0_axis_dim, qnn_indices_bytes)),
|
||||
"QNN does not support negative index values for Gather* ops");
|
||||
} else {
|
||||
gather_indices = std::move(unpacked_tensor);
|
||||
qnn_indices_bytes = std::move(onnx_indices_bytes);
|
||||
}
|
||||
qnn_data_type = QNN_DATATYPE_INT_32;
|
||||
}
|
||||
|
||||
Qnn_TensorType_t tensor_type = qnn_model_wrapper.GetTensorType(input_name);
|
||||
std::vector<uint32_t> input_shape;
|
||||
ORT_RETURN_IF_NOT(qnn_model_wrapper.GetOnnxShape(inputs[1].node_arg, input_shape), "Cannot get shape");
|
||||
std::vector<uint32_t> cast_output_shape(input_shape);
|
||||
QnnTensorWrapper input_tensorwrapper(input_name, tensor_type, qnn_data_type, QnnQuantParamsWrapper(),
|
||||
std::move(input_shape), std::move(gather_indices));
|
||||
std::vector<uint32_t> cast_output_shape(indices_info.shape);
|
||||
QnnTensorWrapper input_tensorwrapper(input_name, tensor_type, indices_info.qnn_data_type, QnnQuantParamsWrapper(),
|
||||
std::move(indices_info.shape), std::move(qnn_indices_bytes));
|
||||
ORT_RETURN_IF_NOT(qnn_model_wrapper.AddTensorWrapper(std::move(input_tensorwrapper)), "Failed to add tensor.");
|
||||
|
||||
if (!is_initializer_input && qnn_data_type == QNN_DATATYPE_INT_64) {
|
||||
// Insert cast node int64 -> int32
|
||||
if (qnn_data_type == QNN_DATATYPE_INT_64) {
|
||||
// Add Cast node for indices
|
||||
indices_input_name = input_name + "_ort_qnn_ep_cast";
|
||||
QnnTensorWrapper cast_output(indices_input_name, QNN_TENSOR_TYPE_NATIVE, QNN_DATATYPE_INT_32,
|
||||
QnnQuantParamsWrapper(), std::move(cast_output_shape));
|
||||
ORT_RETURN_IF_NOT(qnn_model_wrapper.AddTensorWrapper(std::move(cast_output)), "Failed to add tensor.");
|
||||
ORT_RETURN_IF_NOT(qnn_model_wrapper.CreateQnnNode(indices_input_name,
|
||||
QNN_OP_PACKAGE_NAME_QTI_AISW,
|
||||
"Cast",
|
||||
{input_name},
|
||||
{indices_input_name},
|
||||
{},
|
||||
do_op_validation),
|
||||
"Failed to add node.");
|
||||
}
|
||||
// Insert QNN Cast op to convert dynamic indices from int64 to int32.
|
||||
std::string indices_input_name(input_name);
|
||||
if (indices_info.qnn_data_type == QNN_DATATYPE_INT_64) {
|
||||
assert(!indices_info.is_initializer);
|
||||
|
||||
indices_input_name = input_name + "_ort_qnn_ep_cast";
|
||||
QnnTensorWrapper cast_output(indices_input_name, QNN_TENSOR_TYPE_NATIVE, QNN_DATATYPE_INT_32,
|
||||
QnnQuantParamsWrapper(), std::move(cast_output_shape));
|
||||
ORT_RETURN_IF_NOT(qnn_model_wrapper.AddTensorWrapper(std::move(cast_output)), "Failed to add tensor.");
|
||||
ORT_RETURN_IF_NOT(qnn_model_wrapper.CreateQnnNode(indices_input_name,
|
||||
QNN_OP_PACKAGE_NAME_QTI_AISW,
|
||||
"Cast",
|
||||
{input_name},
|
||||
{indices_input_name},
|
||||
{},
|
||||
do_op_validation),
|
||||
"Failed to add node.");
|
||||
}
|
||||
|
||||
input_names.push_back(indices_input_name);
|
||||
|
|
@ -112,20 +193,44 @@ Status GatherOpBuilder::ProcessInputs(QnnModelWrapper& qnn_model_wrapper,
|
|||
return Status::OK();
|
||||
}
|
||||
|
||||
Status GatherOpBuilder::ProcessInputs(QnnModelWrapper& qnn_model_wrapper,
|
||||
const NodeUnit& node_unit,
|
||||
const logging::Logger& logger,
|
||||
std::vector<std::string>& input_names,
|
||||
bool do_op_validation) const {
|
||||
const auto& inputs = node_unit.Inputs();
|
||||
ORT_RETURN_IF(inputs.size() != 2, "QNN EP: ", node_unit.OpType(), " operator must have two inputs");
|
||||
ORT_RETURN_IF_ERROR(ProcessInput(qnn_model_wrapper, inputs[0], logger, input_names));
|
||||
|
||||
int64_t input0_axis_dim = 0;
|
||||
ORT_RETURN_IF_ERROR(GetInpu0AxisDimValue(qnn_model_wrapper, node_unit, /*default_axis*/ 0, input0_axis_dim));
|
||||
|
||||
return ProcessIndicesInput(qnn_model_wrapper, inputs[1], input0_axis_dim, logger, input_names, do_op_validation);
|
||||
}
|
||||
|
||||
Status GatherOpBuilder::ProcessAttributesAndOutputs(QnnModelWrapper& qnn_model_wrapper,
|
||||
const NodeUnit& node_unit,
|
||||
std::vector<std::string>&& input_names,
|
||||
const logging::Logger& logger,
|
||||
bool do_op_validation) const {
|
||||
ORT_UNUSED_PARAMETER(logger);
|
||||
const bool is_gather_elems = node_unit.OpType() == "GatherElements";
|
||||
|
||||
// Create QNN 'axis' parameter.
|
||||
std::vector<std::string> param_tensor_names;
|
||||
int32_t axis_value = 0;
|
||||
Qnn_Scalar_t axis_qnn_scalar = QNN_SCALAR_INIT;
|
||||
ORT_RETURN_IF_ERROR(ProcessAxisAttribute(qnn_model_wrapper, node_unit, axis_qnn_scalar, axis_value));
|
||||
QnnParamWrapper axis_param(node_unit.Index(), node_unit.Name(), QNN_OP_GATHER_PARAM_AXIS, axis_qnn_scalar);
|
||||
QnnParamWrapper axis_param(node_unit.Index(), node_unit.Name(),
|
||||
(is_gather_elems ? QNN_OP_GATHER_ELEMENTS_PARAM_AXIS : QNN_OP_GATHER_PARAM_AXIS),
|
||||
axis_qnn_scalar);
|
||||
param_tensor_names.push_back(axis_param.GetParamTensorName());
|
||||
qnn_model_wrapper.AddParamWrapper(std::move(axis_param));
|
||||
|
||||
if (is_gather_elems) {
|
||||
return ProcessOutputs(qnn_model_wrapper, node_unit, std::move(input_names), std::move(param_tensor_names),
|
||||
logger, do_op_validation, GetQnnOpType(node_unit.OpType()));
|
||||
}
|
||||
|
||||
// if indicies is scalar shape, then need to add Reshape node
|
||||
const auto& input_tensor_wrapper = qnn_model_wrapper.GetQnnTensorWrapper(input_names[0]);
|
||||
const auto& indices_input_tensor_wrapper = qnn_model_wrapper.GetQnnTensorWrapper(input_names[1]);
|
||||
|
|
|
|||
|
|
@ -55,7 +55,7 @@ class QDQGather(QDQOperatorBase):
|
|||
|
||||
def quantize(self):
|
||||
node = self.node
|
||||
assert node.op_type == "Gather"
|
||||
assert node.op_type == "Gather" or node.op_type == "GatherElements"
|
||||
|
||||
if self.quantizer.is_valid_quantize_weight(node.input[0]) or self.quantizer.force_quantize_no_input_check:
|
||||
self.quantizer.quantize_activation_tensor(node.input[0])
|
||||
|
|
|
|||
|
|
@ -79,6 +79,7 @@ QDQRegistry = {
|
|||
"MatMul": QDQMatMul,
|
||||
"Split": QDQSplit,
|
||||
"Gather": QDQGather,
|
||||
"GatherElements": QDQGather,
|
||||
"Where": QDQWhere,
|
||||
"InstanceNormalization": QDQNormalization,
|
||||
"LayerNormalization": QDQNormalization,
|
||||
|
|
|
|||
|
|
@ -827,7 +827,8 @@ select from 'TF8', 'TF16', 'UINT8', 'FLOAT', 'ITENSOR'. \n)");
|
|||
ORT_TSTR("sce_NCd1d2d3_sum_weight_high_ii_expanded"),
|
||||
ORT_TSTR("sce_none_weights_log_prob_expanded"),
|
||||
ORT_TSTR("sce_none_weights_expanded"),
|
||||
ORT_TSTR("convtranspose_3d")};
|
||||
ORT_TSTR("convtranspose_3d"),
|
||||
ORT_TSTR("gather_elements_negative_indices")};
|
||||
|
||||
std::unordered_set<std::basic_string<ORTCHAR_T>> all_disabled_tests(std::begin(immutable_broken_tests), std::end(immutable_broken_tests));
|
||||
|
||||
|
|
|
|||
|
|
@ -42,9 +42,15 @@ void GetData(const std::vector<int64_t>& input_dims, const std::vector<int64_t>&
|
|||
output_data.resize(output_size);
|
||||
std::srand(static_cast<unsigned>(std::time(0)));
|
||||
for (size_t i = 0; i < indices_size; ++i) {
|
||||
#if defined(USE_QNN)
|
||||
// Negative index not possible.
|
||||
indices_data[i] =
|
||||
static_cast<TIndex>(static_cast<int64_t>(std::rand()) % input_dims[axis]);
|
||||
#else
|
||||
// Negative index possible.
|
||||
indices_data[i] =
|
||||
static_cast<TIndex>((static_cast<int64_t>(std::rand()) % (input_dims[axis] * 2)) - input_dims[axis]);
|
||||
#endif
|
||||
}
|
||||
for (size_t i = 0; i < output_size; ++i) {
|
||||
int64_t input_offset = 0;
|
||||
|
|
@ -382,9 +388,10 @@ TEST(GatherElementsOpTest, IndicesOutOfBounds) {
|
|||
// skip cuda as the cuda kernel won't throw the error message
|
||||
// skip openvino which will not throw error message but will ensure no out-of-bound access
|
||||
// skip TensorRT because it doesn't support out of bounds indices
|
||||
// skip QNN because it doesn't support out of bounds indices
|
||||
test.Run(OpTester::ExpectResult::kExpectFailure, "",
|
||||
{kCudaExecutionProvider, kCudaNHWCExecutionProvider, kRocmExecutionProvider, kOpenVINOExecutionProvider,
|
||||
kTensorrtExecutionProvider, kDmlExecutionProvider});
|
||||
kTensorrtExecutionProvider, kDmlExecutionProvider, kQnnExecutionProvider});
|
||||
}
|
||||
|
||||
TEST(GatherElementsOpTest, BigIndices) {
|
||||
|
|
|
|||
319
onnxruntime/test/providers/qnn/gather_elems_op_test.cc
Normal file
319
onnxruntime/test/providers/qnn/gather_elems_op_test.cc
Normal file
|
|
@ -0,0 +1,319 @@
|
|||
// Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
// Licensed under the MIT License.
|
||||
|
||||
#if !defined(ORT_MINIMAL_BUILD)
|
||||
|
||||
#include <optional>
|
||||
#include <string>
|
||||
#include <unordered_map>
|
||||
|
||||
#include "core/graph/node_attr_utils.h"
|
||||
#include "test/optimizer/qdq_test_utils.h"
|
||||
#include "test/providers/qnn/qnn_test_utils.h"
|
||||
|
||||
#include "onnx/onnx_pb.h"
|
||||
|
||||
#include "gtest/gtest.h"
|
||||
|
||||
namespace onnxruntime {
|
||||
namespace test {
|
||||
|
||||
// Creates a graph with a single Q/DQ GatherElements operator. Used for testing HTP backend.
|
||||
template <typename QuantType, typename IndexType>
|
||||
static GetTestQDQModelFn<QuantType> BuildQDQGatherElemsTestCase(const TestInputDef<float>& input_def,
|
||||
const TestInputDef<IndexType>& indices_def,
|
||||
const std::vector<ONNX_NAMESPACE::AttributeProto>& attrs,
|
||||
bool use_contrib_qdq = false) {
|
||||
return [input_def, indices_def, attrs, use_contrib_qdq](ModelTestBuilder& builder,
|
||||
std::vector<QuantParams<QuantType>>& output_qparams) {
|
||||
// input -> Q -> DQ ->
|
||||
NodeArg* input = MakeTestInput(builder, input_def);
|
||||
QuantParams<QuantType> input_qparams = GetTestInputQuantParams<QuantType>(input_def);
|
||||
NodeArg* input_qdq = AddQDQNodePair<QuantType>(builder, input, input_qparams.scale, input_qparams.zero_point,
|
||||
use_contrib_qdq);
|
||||
|
||||
// indices input
|
||||
NodeArg* indices_input = MakeTestInput(builder, indices_def);
|
||||
|
||||
// GatherElements op
|
||||
NodeArg* gather_output = builder.MakeIntermediate();
|
||||
Node& gather_node = builder.AddNode("GatherElements", {input_qdq, indices_input}, {gather_output});
|
||||
|
||||
for (const auto& attr : attrs) {
|
||||
gather_node.AddAttributeProto(attr);
|
||||
}
|
||||
|
||||
// op_output -> Q -> DQ -> output
|
||||
// NOTE: Input and output quantization parameters must be equal for GatherElements.
|
||||
output_qparams[0] = input_qparams; // Overwrite!
|
||||
AddQDQNodePairWithOutputAsGraphOutput<QuantType>(builder, gather_output, input_qparams.scale,
|
||||
input_qparams.zero_point, use_contrib_qdq);
|
||||
};
|
||||
}
|
||||
|
||||
// Runs an GatherElements model on the QNN CPU backend. Checks the graph node assignment, and that inference
|
||||
// outputs for QNN EP and CPU EP match.
|
||||
template <typename DataType, typename IndexType>
|
||||
static void RunCPUGatherElemsOpTest(const TestInputDef<float>& input_def,
|
||||
const TestInputDef<IndexType>& indices_def,
|
||||
const std::vector<ONNX_NAMESPACE::AttributeProto>& attrs,
|
||||
ExpectedEPNodeAssignment expected_ep_assignment,
|
||||
int opset = 13) {
|
||||
ProviderOptions provider_options;
|
||||
float fp32_abs_err = 1e-5f; // default tolerance
|
||||
|
||||
#if defined(_WIN32)
|
||||
provider_options["backend_path"] = "QnnCpu.dll";
|
||||
#else
|
||||
provider_options["backend_path"] = "libQnnCpu.so";
|
||||
#endif
|
||||
|
||||
RunQnnModelTest(BuildOpTestCase<DataType, IndexType>("GatherElements", {input_def}, {indices_def}, attrs),
|
||||
provider_options,
|
||||
opset,
|
||||
expected_ep_assignment,
|
||||
fp32_abs_err);
|
||||
}
|
||||
|
||||
// Runs a QDQ GatherElements model on the QNN HTP backend. Checks the graph node assignment, and that inference
|
||||
// outputs for QNN EP and CPU EP match with expected accuracy.
|
||||
template <typename QuantType, typename IndexType>
|
||||
static void RunHTPQDQGatherElemsOpTest(const TestInputDef<float>& input_def,
|
||||
const TestInputDef<IndexType>& indices_def,
|
||||
const std::vector<ONNX_NAMESPACE::AttributeProto>& attrs,
|
||||
ExpectedEPNodeAssignment expected_ep_assignment,
|
||||
int opset = 13,
|
||||
bool use_contrib_qdq = false) {
|
||||
ProviderOptions provider_options;
|
||||
|
||||
#if defined(_WIN32)
|
||||
provider_options["backend_path"] = "QnnHtp.dll";
|
||||
#else
|
||||
provider_options["backend_path"] = "libQnnHtp.so";
|
||||
#endif
|
||||
|
||||
auto f32_model_builder = BuildOpTestCase<float, IndexType>("GatherElements", {input_def}, {indices_def}, attrs);
|
||||
auto qdq_model_builder = BuildQDQGatherElemsTestCase<QuantType, IndexType>(input_def, indices_def, attrs,
|
||||
use_contrib_qdq);
|
||||
|
||||
TestQDQModelAccuracy<QuantType>(f32_model_builder,
|
||||
qdq_model_builder,
|
||||
provider_options,
|
||||
opset,
|
||||
expected_ep_assignment);
|
||||
}
|
||||
|
||||
// Runs a non-quantized GatherElements model on the QNN HTP backend. Checks the graph node assignment,
|
||||
// and that inference outputs for QNN EP and CPU EP match.
|
||||
template <typename DataType, typename IndexType>
|
||||
static void RunHTPGatherElemsOpTest(const TestInputDef<DataType>& input_def,
|
||||
const TestInputDef<IndexType>& indices_def,
|
||||
const std::vector<ONNX_NAMESPACE::AttributeProto>& attrs,
|
||||
ExpectedEPNodeAssignment expected_ep_assignment,
|
||||
int opset = 13) {
|
||||
ProviderOptions provider_options;
|
||||
float fp32_abs_err = 1e-5f; // default tolerance
|
||||
|
||||
#if defined(_WIN32)
|
||||
provider_options["backend_path"] = "QnnHtp.dll";
|
||||
#else
|
||||
provider_options["backend_path"] = "libQnnHtp.so";
|
||||
#endif
|
||||
|
||||
RunQnnModelTest(BuildOpTestCase<DataType, IndexType>("GatherElements", {input_def}, {indices_def}, attrs),
|
||||
provider_options,
|
||||
opset,
|
||||
expected_ep_assignment,
|
||||
fp32_abs_err);
|
||||
}
|
||||
|
||||
//
|
||||
// CPU tests:
|
||||
//
|
||||
|
||||
// Test GatherElements op on CPU backend:
|
||||
// positive, dynamic, int64 indices.
|
||||
TEST_F(QnnCPUBackendTests, GatherElems_DataF32_IndicesInt64) {
|
||||
RunCPUGatherElemsOpTest<float, int64_t>(
|
||||
TestInputDef<float>({3, 3}, false, {1.f, 2.f, 3.f, 4.f, 5.f, 6.f, 7.f, 8.f, 9.f}),
|
||||
TestInputDef<int64_t>({2, 3}, false, {1, 2, 0, 2, 0, 0}),
|
||||
{utils::MakeAttribute("axis", static_cast<int64_t>(1))},
|
||||
ExpectedEPNodeAssignment::All);
|
||||
}
|
||||
|
||||
// Test GatherElements op on CPU backend:
|
||||
// positive, dynamic, int32 indices.
|
||||
TEST_F(QnnCPUBackendTests, GatherElems_DataF32_IndicesInt32) {
|
||||
RunCPUGatherElemsOpTest<float, int32_t>(
|
||||
TestInputDef<float>({3, 3}, false, {1.f, 2.f, 3.f, 4.f, 5.f, 6.f, 7.f, 8.f, 9.f}),
|
||||
TestInputDef<int32_t>({2, 3}, false, {1, 2, 0, 2, 0, 0}),
|
||||
{utils::MakeAttribute("axis", static_cast<int64_t>(1))},
|
||||
ExpectedEPNodeAssignment::All);
|
||||
}
|
||||
|
||||
// Test GatherElements op on CPU backend:
|
||||
// positive, static, int64 indices.
|
||||
TEST_F(QnnCPUBackendTests, GatherElems_DataF32_StaticIndicesInt64) {
|
||||
RunCPUGatherElemsOpTest<float, int64_t>(
|
||||
TestInputDef<float>({3, 3}, false, {1.f, 2.f, 3.f, 4.f, 5.f, 6.f, 7.f, 8.f, 9.f}),
|
||||
TestInputDef<int64_t>({2, 3}, true, {1, 2, 0, 2, 0, 0}),
|
||||
{utils::MakeAttribute("axis", static_cast<int64_t>(1))},
|
||||
ExpectedEPNodeAssignment::All);
|
||||
}
|
||||
|
||||
// Test GatherElements op on CPU backend:
|
||||
// positive, static, int32 indices.
|
||||
TEST_F(QnnCPUBackendTests, GatherElems_DataF32_StaticIndicesInt32) {
|
||||
RunCPUGatherElemsOpTest<float, int32_t>(
|
||||
TestInputDef<float>({3, 3}, false, {1.f, 2.f, 3.f, 4.f, 5.f, 6.f, 7.f, 8.f, 9.f}),
|
||||
TestInputDef<int32_t>({2, 3}, true, {1, 2, 0, 2, 0, 0}),
|
||||
{utils::MakeAttribute("axis", static_cast<int64_t>(1))},
|
||||
ExpectedEPNodeAssignment::All);
|
||||
}
|
||||
|
||||
#if defined(__aarch64__) || defined(_M_ARM64) || defined(__linux__)
|
||||
//
|
||||
// HTP tests:
|
||||
//
|
||||
|
||||
// Test non-quantized GatherElements op on HTP backend:
|
||||
// Input[0] is int32. Indices are int32, positive, and dynamic. Both inputs are rank 1.
|
||||
TEST_F(QnnHTPBackendTests, GatherElems_DataInt32_IndicesInt32_Rank1) {
|
||||
RunHTPGatherElemsOpTest<int32_t, int32_t>(
|
||||
TestInputDef<int32_t>({3}, false, {1, 2, 3}),
|
||||
TestInputDef<int32_t>({2}, false, {1, 2}),
|
||||
{utils::MakeAttribute("axis", static_cast<int64_t>(0))},
|
||||
ExpectedEPNodeAssignment::All);
|
||||
}
|
||||
|
||||
// Test uint8 QDQ GatherElements op on HTP backend:
|
||||
// positive, dynamic, int32 indices.
|
||||
TEST_F(QnnHTPBackendTests, GatherElems_DataUint8_IndicesInt32) {
|
||||
RunHTPQDQGatherElemsOpTest<uint8_t, int32_t>(
|
||||
TestInputDef<float>({3, 3}, false, {1.f, 2.f, 3.f, 4.f, 5.f, 6.f, 7.f, 8.f, 9.f}),
|
||||
TestInputDef<int32_t>({2, 3}, false, {1, 2, 0, 2, 0, 0}),
|
||||
{utils::MakeAttribute("axis", static_cast<int64_t>(1))},
|
||||
ExpectedEPNodeAssignment::All);
|
||||
}
|
||||
|
||||
// Test uint8 QDQ GatherElements op on HTP backend:
|
||||
// positive, static, int32 indices.
|
||||
TEST_F(QnnHTPBackendTests, GatherElems_DataUint8_StaticIndicesInt32) {
|
||||
RunHTPQDQGatherElemsOpTest<uint8_t, int32_t>(
|
||||
TestInputDef<float>({3, 3}, false, {1.f, 2.f, 3.f, 4.f, 5.f, 6.f, 7.f, 8.f, 9.f}),
|
||||
TestInputDef<int32_t>({2, 3}, true, {1, 2, 0, 2, 0, 0}),
|
||||
{utils::MakeAttribute("axis", static_cast<int64_t>(1))},
|
||||
ExpectedEPNodeAssignment::All);
|
||||
}
|
||||
|
||||
// Test GatherElements op on HTP backend:
|
||||
// positive, static, int64 indices.
|
||||
// HTP does not support int64, but QNN EP converts static int64 indices into int32.
|
||||
TEST_F(QnnHTPBackendTests, GatherElems_DataUint8_StaticIndicesInt64) {
|
||||
RunHTPQDQGatherElemsOpTest<uint8_t, int64_t>(
|
||||
TestInputDef<float>({3, 3}, false, {1.f, 2.f, 3.f, 4.f, 5.f, 6.f, 7.f, 8.f, 9.f}),
|
||||
TestInputDef<int64_t>({2, 3}, true, {1, 2, 0, 2, 0, 0}),
|
||||
{utils::MakeAttribute("axis", static_cast<int64_t>(1))},
|
||||
ExpectedEPNodeAssignment::All);
|
||||
}
|
||||
|
||||
// Test GatherElements op on HTP backend:
|
||||
// negative, static, int32 indices.
|
||||
TEST_F(QnnHTPBackendTests, GatherElems_DataUint8_StaticNegIndicesInt32) {
|
||||
RunHTPQDQGatherElemsOpTest<uint8_t, int32_t>(
|
||||
TestInputDef<float>({3, 3}, false, {1.f, 2.f, 3.f, 4.f, 5.f, 6.f, 7.f, 8.f, 9.f}),
|
||||
TestInputDef<int32_t>({2, 3}, true, {1, 2, -3, 2, 0, 0}),
|
||||
{utils::MakeAttribute("axis", static_cast<int64_t>(1))},
|
||||
ExpectedEPNodeAssignment::All);
|
||||
}
|
||||
|
||||
// Test GatherElements op on HTP backend:
|
||||
// negative, static, int64 indices.
|
||||
TEST_F(QnnHTPBackendTests, GatherElems_DataUint8_StaticNegIndicesInt64) {
|
||||
RunHTPQDQGatherElemsOpTest<uint8_t, int64_t>(
|
||||
TestInputDef<float>({3, 3}, false, {1.f, 2.f, 3.f, 4.f, 5.f, 6.f, 7.f, 8.f, 9.f}),
|
||||
TestInputDef<int64_t>({2, 3}, true, {1, -1, -3, 2, 0, 0}),
|
||||
{utils::MakeAttribute("axis", static_cast<int64_t>(1))},
|
||||
ExpectedEPNodeAssignment::All);
|
||||
}
|
||||
|
||||
// Test QDQ GatherElements op on HTP backend:
|
||||
// Input[0] is uint16_t, and rank 4. Indices are int64_t, rank 4, negative, and static.
|
||||
// Axis is negative (points to last dim).
|
||||
TEST_F(QnnHTPBackendTests, GatherElems_DataUint16_StaticNegIndicesInt64) {
|
||||
const std::vector<int64_t> input_shape = {1, 2, 3, 3};
|
||||
const std::vector<float> input_data = GetSequentialFloatData(input_shape, -8.0f, 1.0f);
|
||||
RunHTPQDQGatherElemsOpTest<uint16_t, int64_t>(
|
||||
TestInputDef<float>(input_shape, false, input_data),
|
||||
TestInputDef<int64_t>({1, 1, 2, 2}, true, {0, -1, -3, 2}),
|
||||
{utils::MakeAttribute("axis", static_cast<int64_t>(-1))},
|
||||
ExpectedEPNodeAssignment::All,
|
||||
/*opset*/ 21);
|
||||
}
|
||||
|
||||
// Test QDQ GatherElements op on HTP backend with large number of indices:
|
||||
TEST_F(QnnHTPBackendTests, GatherElems_DataUint16_StaticNegIndicesInt64_Large) {
|
||||
const std::vector<int64_t> input_shape = {12, 1024, 512};
|
||||
std::vector<float> input_data(12 * 1024 * 512);
|
||||
for (size_t i = 0; i < input_data.size(); i++) {
|
||||
input_data[i] = static_cast<float>((static_cast<int64_t>(i) % 8));
|
||||
}
|
||||
|
||||
RunHTPQDQGatherElemsOpTest<uint16_t, int64_t>(
|
||||
TestInputDef<float>(input_shape, false, input_data),
|
||||
TestInputDef<int64_t>({12, 1024, 1024}, true, -512, 511),
|
||||
{utils::MakeAttribute("axis", static_cast<int64_t>(-1))},
|
||||
ExpectedEPNodeAssignment::All,
|
||||
/*opset*/ 21);
|
||||
}
|
||||
|
||||
// Test QDQ GatherElements op on HTP backend with large number of indices.
|
||||
// TODO: Investigate inaccuracy.
|
||||
// Negative input[0] values seem to cause inaccuracies with 2M+ indices.
|
||||
// Inaccuracy detected for output 'output_0', element 131072
|
||||
// output_range=300, tolerance=0.40000000596046448%.
|
||||
// Expected val (f32@CPU_EP): 121
|
||||
// qdq@QNN_EP val: -97.999542236328125 (err: 218.99954223632812, err/output_range: 72.999847412109375%)
|
||||
// qdq@CPU_EP val: 120.99794006347656 (err: 0.0020599365234375, err/output_range: 0.0006866455078125%)
|
||||
// abs(qdq@QNN_EP - qdq@CPU_EP) / output_range = 72.999160766601562%
|
||||
TEST_F(QnnHTPBackendTests, DISABLED_GatherElems_DataUint16_StaticNegIndicesInt64_Large2) {
|
||||
// Input data with sequential values from -98.0f to 202.0f
|
||||
const std::vector<int64_t> input_shape = {12, 1024, 512};
|
||||
std::vector<float> input_data(12 * 1024 * 512);
|
||||
for (size_t i = 0; i < input_data.size(); i++) {
|
||||
int32_t int_val = -98 + (static_cast<int32_t>(i) % 301);
|
||||
input_data[i] = static_cast<float>(int_val);
|
||||
}
|
||||
|
||||
// Indices with values between -512 to 511.
|
||||
const std::vector<int64_t> indices_shape = {12, 1024, 1024};
|
||||
std::vector<int64_t> indices(12 * 1024 * 1024);
|
||||
for (size_t i = 0; i < indices.size(); i++) {
|
||||
indices[i] = static_cast<int64_t>(-512 + (static_cast<int32_t>(i) % 1024));
|
||||
}
|
||||
|
||||
RunHTPQDQGatherElemsOpTest<uint16_t, int64_t>(
|
||||
TestInputDef<float>(input_shape, false, input_data),
|
||||
TestInputDef<int64_t>(indices_shape, true, indices),
|
||||
{utils::MakeAttribute("axis", static_cast<int64_t>(-1))},
|
||||
ExpectedEPNodeAssignment::All,
|
||||
/*opset*/ 21);
|
||||
}
|
||||
|
||||
// Test GatherElements op on HTP backend:
|
||||
// Tests that dynamic int64 indices are supported on HTP backend if the indices are a graph input.
|
||||
// QNN SDK 2.23 added support for Cast from int64 to int32.
|
||||
TEST_F(QnnHTPBackendTests, GatherElems_DynamicInt64IndicesSupportedAsGraphInput) {
|
||||
RunHTPQDQGatherElemsOpTest<uint8_t, int64_t>(
|
||||
TestInputDef<float>({3, 3}, false, {1.f, 2.f, 3.f, 4.f, 5.f, 6.f, 7.f, 8.f, 9.f}),
|
||||
TestInputDef<int64_t>({2, 3}, false, {1, 2, 0, 2, 0, 0}),
|
||||
{utils::MakeAttribute("axis", static_cast<int64_t>(1))},
|
||||
ExpectedEPNodeAssignment::All);
|
||||
}
|
||||
|
||||
#endif // defined(__aarch64__) || defined(_M_ARM64) || defined(__linux__)
|
||||
|
||||
} // namespace test
|
||||
} // namespace onnxruntime
|
||||
|
||||
#endif // !defined(ORT_MINIMAL_BUILD)
|
||||
|
|
@ -97,13 +97,14 @@ TEST_F(QnnHTPBackendTests, GatherOp_U16_IndicesStaticInt64_Axis0) {
|
|||
true); // Use 'com.microsoft' Q/DQ ops
|
||||
}
|
||||
|
||||
// Tests that dynamic int64 indices are not supported on HTP backend.
|
||||
// Tests that dynamic int64 indices are supported on HTP backend if the indices are a graph input.
|
||||
// QNN SDK 2.23 added support for Cast from int64 to int32.
|
||||
TEST_F(QnnHTPBackendTests, GatherOp_IndicesDynamicInt64_Axis0) {
|
||||
RunQDQGatherOpTest<uint8_t, int64_t>(TestInputDef<float>({3, 2}, false, {1.0f, 1.2f, 2.3f, 3.4f, 4.5f, 5.7f}),
|
||||
TestInputDef<int64_t>({2, 2}, false, {0, 1, 1, 2}),
|
||||
{utils::MakeAttribute("axis", static_cast<int64_t>(0))},
|
||||
13,
|
||||
ExpectedEPNodeAssignment::None);
|
||||
ExpectedEPNodeAssignment::All);
|
||||
}
|
||||
|
||||
// Test creates a DQ -> Gather -> Q -> DQ graph, and checks that all
|
||||
|
|
|
|||
110
onnxruntime/test/python/quantization/test_op_gather_elems.py
Normal file
110
onnxruntime/test/python/quantization/test_op_gather_elems.py
Normal file
|
|
@ -0,0 +1,110 @@
|
|||
#!/usr/bin/env python
|
||||
# -------------------------------------------------------------------------
|
||||
# Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
# Licensed under the MIT License. See License.txt in the project root for
|
||||
# license information.
|
||||
# --------------------------------------------------------------------------
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import tempfile
|
||||
import unittest
|
||||
|
||||
import numpy as np
|
||||
import onnx
|
||||
from op_test_utils import TestDataFeeds, check_op_type_count
|
||||
|
||||
from onnxruntime.quantization import QuantFormat, QuantType, quantize_static
|
||||
|
||||
|
||||
class TestQDQGatherElements(unittest.TestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls._tmp_model_dir = tempfile.TemporaryDirectory(prefix="ort.qdq.gatherelems_")
|
||||
|
||||
# Note: swap with the commented line if you want to see the models in local test dir.
|
||||
cls._tmp_dir_path = cls._tmp_model_dir.name
|
||||
# cls._tmp_dir_path = "."
|
||||
|
||||
@classmethod
|
||||
def tearDownClass(cls):
|
||||
cls._tmp_model_dir.cleanup()
|
||||
|
||||
def build_test_model_static_indices(
|
||||
self,
|
||||
inp_shape: list[int],
|
||||
indices_data: np.ndarray,
|
||||
axis: int = 0,
|
||||
):
|
||||
input_0 = onnx.helper.make_tensor_value_info("input_0", onnx.TensorProto.FLOAT, inp_shape)
|
||||
output_0 = onnx.helper.make_tensor_value_info("output_0", onnx.TensorProto.FLOAT, None)
|
||||
indices = onnx.numpy_helper.from_array(indices_data, "indices")
|
||||
|
||||
gatherelems_node = onnx.helper.make_node(
|
||||
"GatherElements", ["input_0", "indices"], ["output_0"], axis=axis, name="GatherElems0"
|
||||
)
|
||||
graph = onnx.helper.make_graph(
|
||||
[gatherelems_node],
|
||||
"GatherElemsf32",
|
||||
[input_0],
|
||||
[output_0],
|
||||
initializer=[indices],
|
||||
)
|
||||
opset_imports = [onnx.helper.make_opsetid("", 21)]
|
||||
model = onnx.helper.make_model(graph, opset_imports=opset_imports)
|
||||
|
||||
return onnx.shape_inference.infer_shapes(model)
|
||||
|
||||
def test_qdq_gatherelems_static_indices(self):
|
||||
"""
|
||||
Test quantization of GatherElements with static indices.
|
||||
"""
|
||||
float_model_path = os.path.join(self._tmp_dir_path, "gather_elems.f32.onnx")
|
||||
qdq_model_path = os.path.join(self._tmp_dir_path, "gather_elems.qdq.onnx")
|
||||
|
||||
inp_shape = [3, 3]
|
||||
indices_data = np.array([[1, 2, 0], [2, 0, 0]], dtype=np.int64)
|
||||
float_model = self.build_test_model_static_indices(inp_shape, indices_data, axis=0)
|
||||
|
||||
onnx.checker.check_model(float_model, True)
|
||||
onnx.save_model(float_model, float_model_path)
|
||||
|
||||
input_data1 = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]], dtype=np.float32)
|
||||
input_data2 = np.array([[8, 2, 9], [4, 5, 6], [7, 8, 1]], dtype=np.float32)
|
||||
data_reader = TestDataFeeds([{"input_0": input_data1}, {"input_0": input_data2}])
|
||||
|
||||
quantize_static(
|
||||
float_model_path,
|
||||
qdq_model_path,
|
||||
data_reader,
|
||||
quant_format=QuantFormat.QDQ,
|
||||
activation_type=QuantType.QUInt16,
|
||||
weight_type=QuantType.QUInt8,
|
||||
op_types_to_quantize=[node.op_type for node in float_model.graph.node],
|
||||
extra_options={
|
||||
"ForceQuantizeNoInputCheck": True,
|
||||
},
|
||||
)
|
||||
|
||||
qdq_node_counts = {"QuantizeLinear": 2, "DequantizeLinear": 2}
|
||||
check_op_type_count(self, qdq_model_path, **qdq_node_counts)
|
||||
|
||||
qdq_model = onnx.load_model(qdq_model_path)
|
||||
onnx.checker.check_model(qdq_model, True)
|
||||
|
||||
initializers = {init.name: init for init in qdq_model.graph.initializer}
|
||||
|
||||
zp_name = "input_0_zero_point"
|
||||
scale_name = "input_0_scale"
|
||||
self.assertIn(zp_name, initializers)
|
||||
self.assertIn(scale_name, initializers)
|
||||
|
||||
# Check that all Q/DQ nodes use the same scale and zero-point
|
||||
for node in qdq_model.graph.node:
|
||||
if node.op_type == "QuantizeLinear" or node.op_type == "DequantizeLinear":
|
||||
self.assertEqual(node.input[1], scale_name)
|
||||
self.assertEqual(node.input[2], zp_name)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Loading…
Reference in a new issue