saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-26 19:52:38 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Fuse DQ -> ArgMax into ArgMax (#10274)

Browse source
This commit is contained in:
Yi-Hong Lyu 2022-01-19 06:47:33 +08:00 committed by GitHub
parent e27f2dc932
commit 62eab67f79
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
6 changed files with 135 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

26
onnxruntime/core/optimizer/qdq_transformer/qdq_util.cc
View file

@ -58,5 +58,31 @@ bool IsQDQPairSupported(
*q_scale.data<float>() == *dq_scale.data<float>();
}
bool IsDQSupported(
const Node& dq_node,
const std::function<const ONNX_NAMESPACE::TensorProto*(const std::string&)>& get_const_initializer) {
ConstPointerContainer<std::vector<NodeArg*>> dq_input_defs = dq_node.InputDefs();
// DQ contains optional input is not supported
// non-scalar DQ scale and zero point needs are not supported
if (dq_input_defs.size() != InputIndex::TOTAL_COUNT ||
!optimizer_utils::IsScalar(*dq_input_defs[InputIndex::SCALE_ID]) ||
!optimizer_utils::IsScalar(*dq_input_defs[InputIndex::ZERO_POINT_ID])) {
return false;
}
// if DQ scale and zero point are not constant, return false
const ONNX_NAMESPACE::TensorProto* dq_scale_tensor_proto =
get_const_initializer(dq_input_defs[InputIndex::SCALE_ID]->Name());
const ONNX_NAMESPACE::TensorProto* dq_zp_tensor_proto =
get_const_initializer(dq_input_defs[InputIndex::ZERO_POINT_ID]->Name());
if (nullptr == dq_zp_tensor_proto ||
nullptr == dq_scale_tensor_proto) {
return false;
}
return true;
}
} // namespace QDQ
} // namespace onnxruntime

7
onnxruntime/core/optimizer/qdq_transformer/qdq_util.h
View file

@ -36,5 +36,12 @@ bool IsQDQPairSupported(
const std::function<const ONNX_NAMESPACE::TensorProto*(const std::string&)>& get_const_initializer,
const Path& model_path);
// Check if DQ is supported in the QDQ transformer. It requires:
// 1. DQ doesn't have optional input.
// 2. scale and zero point is constant scalar
bool IsDQSupported(
const Node& dq_node,
const std::function<const ONNX_NAMESPACE::TensorProto*(const std::string&)>& get_const_initializer);
} // namespace QDQ
} // namespace onnxruntime

24
onnxruntime/core/optimizer/qdq_transformer/selectors_actions/qdq_selector_action_transformer.cc
View file

@ -45,6 +45,29 @@ void DropQDQNodesRules(SelectorActionRegistry& qdq_selector_action_registry) {
#endif
}
// create rules for ops that don't change the data
void DropDQNodesRules(SelectorActionRegistry& qdq_selector_action_registry) {
// 2 nodes. DQ, target. Merge into target and remove DQ.
const std::string action_name{"dropDQ"};
NTO::NodeLocation dq{NTO::NodeType::kInput, 0};
// Move DQ input 0 to target input 0.
std::vector<NodeAndMoveInfo> moves{
MoveToSlot(dq, ArgType::kInput, 0, ArgType::kInput, 0)};
std::unique_ptr<Action> action = std::make_unique<MergeIntoTarget>(std::move(moves));
#if !defined(ORT_MINIMAL_BUILD)
std::unique_ptr<NodeSelector> selector = std::make_unique<QDQ::DropDQNodesSelector>();
qdq_selector_action_registry.RegisterSelectorAndAction(action_name,
{{"ArgMax", {}}},
std::move(selector),
std::move(action));
#else
qdq_selector_action_registry.RegisterAction(action_name, std::move(action));
#endif
}
void UnaryOpQDQRules(SelectorActionRegistry& qdq_selector_action_registry) {
// 3 nodes. DQ, target, Q
// Replace with internal QLinear version of operator. Delete all original nodes.
@ -148,6 +171,7 @@ SelectorActionRegistry CreateSelectorActionRegistry(bool is_int8_allowed) {
SelectorActionRegistry qdq_selector_action_registry;
DropQDQNodesRules(qdq_selector_action_registry);
DropDQNodesRules(qdq_selector_action_registry);
UnaryOpQDQRules(qdq_selector_action_registry);
BinaryOpQDQRules(qdq_selector_action_registry);
VariadicOpQDQRules(qdq_selector_action_registry);

24
onnxruntime/core/optimizer/qdq_transformer/selectors_actions/qdq_selectors.cc
View file

@ -104,6 +104,30 @@ bool DropQDQNodeGroupSelector::Check(const GraphViewer& graph_viewer,
return IsQDQPairSupported(q_node, dq_node, get_const_initializer, graph_viewer.ModelPath());
}
bool DropDQNodeGroupSelector::CheckDQNodes(const Node& node, const std::vector<const Node*>& dq_nodes) const {
int num_dq_inputs = NumActualValues(node, true);
return num_dq_inputs == gsl::narrow_cast<int>(dq_nodes.size());
}
bool DropDQNodeGroupSelector::Check(const GraphViewer& graph_viewer,
const Node& node,
const std::vector<const Node*>& dq_nodes,
const std::vector<const Node*>& q_nodes) const {
if (!CheckDQNodes(node, dq_nodes)) {
return false;
}
(void)q_nodes;
const Node& dq_node = *dq_nodes.front();
auto get_const_initializer = [&graph_viewer](const std::string& initializer_name) {
return graph_viewer.GetConstantInitializer(initializer_name, true);
};
return IsDQSupported(dq_node, get_const_initializer);
}
bool UnaryNodeGroupSelector::Check(const GraphViewer& graph_viewer, const Node& node,
const std::vector<const Node*>& dq_nodes,
const std::vector<const Node*>& q_nodes) const {

15
onnxruntime/core/optimizer/qdq_transformer/selectors_actions/qdq_selectors.h
View file

@ -56,6 +56,16 @@ class DropQDQNodeGroupSelector : public NodeGroupSelector {
const std::vector<const Node*>& q_nodes) const override;
};
// Single DQ -> node.
class DropDQNodeGroupSelector : public NodeGroupSelector {
// base check that we have the expected number of DQ inputs.
bool CheckDQNodes(const Node& node, const std::vector<const Node*>& dq_nodes) const;
bool Check(const GraphViewer& graph_viewer, const Node& node,
const std::vector<const Node*>& dq_nodes,
const std::vector<const Node*>& q_nodes) const override;
};
// single input. default is to only support uint8.
class UnaryNodeGroupSelector : public NodeGroupSelector {
bool Check(const GraphViewer& graph_viewer, const Node& node,
@ -142,6 +152,11 @@ class DropQDQNodesSelector : public BaseSelector {
DropQDQNodesSelector() : BaseSelector(std::make_unique<DropQDQNodeGroupSelector>()) {}
};
class DropDQNodesSelector : public BaseSelector {
public:
DropDQNodesSelector() : BaseSelector(std::make_unique<DropDQNodeGroupSelector>()) {}
};
class UnarySelector : public BaseSelector {
public:
UnarySelector() : BaseSelector(std::make_unique<UnaryNodeGroupSelector>()) {}

39
onnxruntime/test/optimizer/qdq_transformer_test.cc
View file

@ -819,6 +819,45 @@ TEST(QDQTransformerTests, ResizeReshape) {
test_case({1, 2, 26, 42}, {4});
}
TEST(QDQTransformerTests, ArgMax) {
auto test_case = [&](const std::vector<int64_t>& input_shape,
int axis,
int keepdims,
int select_last_index) {
auto build_test_case = [&](ModelTestBuilder& builder) {
auto* input_arg = builder.MakeInput<uint8_t>(input_shape,
std::numeric_limits<uint8_t>::min(),
std::numeric_limits<uint8_t>::max());
auto* output_arg = builder.MakeOutput();
// add DQ
auto* dq_output = builder.MakeIntermediate();
builder.AddDequantizeLinearNode<uint8_t>(input_arg, .003f, 1, dq_output);
// add ArgMax
Node& argmax_node = builder.AddNode("ArgMax", {dq_output}, {output_arg});
argmax_node.AddAttribute("axis", static_cast<int64_t>(axis));
argmax_node.AddAttribute("keepdims", static_cast<int64_t>(keepdims));
argmax_node.AddAttribute("select_last_index", static_cast<int64_t>(select_last_index));
};
auto check_argmax_graph = [&](InferenceSessionWrapper& session) {
auto op_to_count = CountOpsInGraph(session.GetGraph());
EXPECT_EQ(op_to_count["ArgMax"], 1);
EXPECT_EQ(op_to_count["DequantizeLinear"], 0);
};
TransformerTester(build_test_case, check_argmax_graph,
TransformerLevel::Level1,
TransformerLevel::Level2,
/* opset_version */ 13);
};
test_case({2, 13, 12, 37}, 1, 0, 0);
test_case({2, 13, 12, 37}, 0, 1, 0);
test_case({2, 13, 12, 37}, 0, 0, 1);
}
TEST(QDQTransformerTests, QLinearMatMul) {
auto test_case = [&](const std::vector<int64_t>& input1_shape, const std::vector<int64_t>& input2_shape) {
auto build_test_case = [&](ModelTestBuilder& builder) {