From 46246f1bbdb518454fed997087576da8367f220f Mon Sep 17 00:00:00 2001 From: Zhang Lei Date: Thu, 13 May 2021 09:50:37 -0700 Subject: [PATCH] Add nhwc transformer support and unittest for qlinear concat. (#7587) --- .../core/optimizer/nhwc_transformer.cc | 72 ++++++++++++++++--- .../optimizer/graph_transform_test_builder.h | 16 +++++ .../test/optimizer/nhwc_transformer_test.cc | 44 ++++++++++++ 3 files changed, 123 insertions(+), 9 deletions(-) diff --git a/onnxruntime/core/optimizer/nhwc_transformer.cc b/onnxruntime/core/optimizer/nhwc_transformer.cc index 3691fdf882..26829f71fc 100644 --- a/onnxruntime/core/optimizer/nhwc_transformer.cc +++ b/onnxruntime/core/optimizer/nhwc_transformer.cc @@ -40,6 +40,25 @@ class NhwcTransformerImpl { return (it != nhwc_args_.end()) ? it->second.get() : nullptr; } + static bool NchwAxisToNhwc(int64_t& axis, int rank) { + if (axis < -rank || axis >= rank) { + return false; + } + bool is_negative_axis = (axis < 0); + if (is_negative_axis) { + axis = axis + rank; + } + if (axis == 1) { + axis = rank - 1; + } else if (axis > 1) { + axis = axis - 1; + } + if (is_negative_axis) { + axis = axis - rank; + } + return true; + } + size_t RemoveOutputEdge(Node& node, size_t output_index); void CreateNhwcArgument(Node& node, Node& nhwc_node, int rank, size_t output_index); void CreateNhwcArgument(Node& node, Node& nhwc_node, int rank); @@ -52,6 +71,7 @@ class NhwcTransformerImpl { void TransformMaxPool(Node& node); void TransformSplit(Node& node); void TransformPad(Node& node); + void TransformQLinearConcat(Node& node); Graph& graph_; @@ -228,6 +248,45 @@ void NhwcTransformerImpl::TransformQLinearActivation(Node& node) { CreateNhwcArgument(node, node, nhwc_input->rank_); } +void NhwcTransformerImpl::TransformQLinearConcat(Node& node) { + auto& input_defs = node.MutableInputDefs(); + + int rank = 0; + for (size_t def_index = 2, def_count = input_defs.size(); def_index < def_count; def_index += 3) { + auto* nhwc_input = LookupNhwcArgument(input_defs[def_index]); + if (nhwc_input == nullptr || (def_index > 2 && nhwc_input->rank_ != rank)) { + return; + } + if (def_index == 2) { + rank = nhwc_input->rank_; + } + } + + // Change the axis attribute accordingly for NCHW to NHWC model. + const auto* axis_attr = graph_utils::GetNodeAttribute(node, "axis"); + if (axis_attr != nullptr && utils::HasInt(*axis_attr)) { + int64_t axis = axis_attr->i(); + if (!NchwAxisToNhwc(axis, rank)) { + // direct return on invalid axis + return; + } + node.AddAttribute("axis", axis); + } else { + // direct return on invalid node + return; + } + + for (size_t def_index = 2, def_count = input_defs.size(); def_index < def_count; def_index += 3) { + // Update the node to directly use the NHWC inputs and decrement the original + // use counts of the NHWC inputs. + auto* nhwc_input = LookupNhwcArgument(input_defs[def_index]); + input_defs[def_index] = nhwc_input->nhwc_arg_; + nhwc_input->remaining_original_uses_--; + } + + CreateNhwcArgument(node, node, rank); +} + void NhwcTransformerImpl::TransformQLinearGlobalAveragePool(Node& node) { auto& input_defs = node.MutableInputDefs(); @@ -300,20 +359,13 @@ void NhwcTransformerImpl::TransformSplit(Node& node) { const auto* axis_attr = graph_utils::GetNodeAttribute(node, "axis"); if (axis_attr != nullptr && utils::HasInt(*axis_attr)) { int64_t axis = axis_attr->i(); - if (axis < -nhwc_input->rank_ || axis >= nhwc_input->rank_) { + if (!NchwAxisToNhwc(axis, nhwc_input->rank_)) { // direct return on invalid axis return; } - if (axis < 0) { - axis = axis + nhwc_input->rank_; - } - if (axis == 1) { - axis = nhwc_input->rank_ - 1; - } else if (axis > 1) { - axis = axis - 1; - } node.AddAttribute("axis", axis); } + // default axis is 0 when attribute not exists, which do not need update // Update the node to directly use the NHWC inputs and decrement the original // use counts of the NHWC inputs. @@ -377,6 +429,8 @@ void NhwcTransformerImpl::Transform(Node& node) { TransformQLinearActivation(node); } else if (graph_utils::IsSupportedOptypeVersionAndDomain(node, "QLinearGlobalAveragePool", {1}, kMSDomain)) { TransformQLinearGlobalAveragePool(node); + } else if (graph_utils::IsSupportedOptypeVersionAndDomain(node, "QLinearConcat", {1}, kMSDomain)) { + TransformQLinearConcat(node); } else if (graph_utils::IsSupportedOptypeVersionAndDomain(node, "MaxPool", {12})) { TransformMaxPool(node); } else if (graph_utils::IsSupportedOptypeVersionAndDomain(node, "Split", {2, 11, 13})) { diff --git a/onnxruntime/test/optimizer/graph_transform_test_builder.h b/onnxruntime/test/optimizer/graph_transform_test_builder.h index e9d7453402..c5f403161a 100644 --- a/onnxruntime/test/optimizer/graph_transform_test_builder.h +++ b/onnxruntime/test/optimizer/graph_transform_test_builder.h @@ -213,6 +213,22 @@ class ModelTestBuilder { return AddNode(op_type, input_args, {output_arg}, kMSDomain); } + Node& AddQLinearConcatLike(const std::string& op_type, + NodeArg* output_arg, + float output_scale, + uint8_t output_zero_point, + std::vector> quantized_inputs) { + std::vector input_args; + input_args.push_back(MakeScalarInitializer(output_scale)); + input_args.push_back(MakeScalarInitializer(output_zero_point)); + for (size_t input_index = 0; input_index < quantized_inputs.size(); ++input_index) { + input_args.push_back(std::get<0>(quantized_inputs[input_index])); + input_args.push_back(MakeScalarInitializer(std::get<1>(quantized_inputs[input_index]))); + input_args.push_back(MakeScalarInitializer(std::get<2>(quantized_inputs[input_index]))); + } + return AddNode(op_type, input_args, {output_arg}, kMSDomain); + } + Node& AddQLinearActivationNode(const std::string& op_type, NodeArg* input_arg, float input_scale, diff --git a/onnxruntime/test/optimizer/nhwc_transformer_test.cc b/onnxruntime/test/optimizer/nhwc_transformer_test.cc index 474b6f1c7e..55d3dd2f4b 100644 --- a/onnxruntime/test/optimizer/nhwc_transformer_test.cc +++ b/onnxruntime/test/optimizer/nhwc_transformer_test.cc @@ -290,6 +290,50 @@ TEST(NhwcTransformerTests, ConvSplit) { } } +TEST(NhwcTransformerTests, ConvSplitQLinearConcat) { + for (int64_t axis = -4LL; axis < 4; axis++) { + auto build_test_case = [&, axis](ModelTestBuilder& builder) { + auto* input_arg = builder.MakeInput({2, 23, 16, 16}, 0, 31); + auto* conv_output_arg = builder.MakeIntermediate(); + auto* split_output1_arg = builder.MakeIntermediate(); + auto* split_output2_arg = builder.MakeIntermediate(); + auto* qlconcat_output_arg = builder.MakeIntermediate(); + auto* output_arg = builder.MakeOutput(); + + const int64_t conv1_output_channels = 32; + auto* conv1_weight_arg = NhwcMakeInitializer(builder, {conv1_output_channels, 23, 3, 3}); + Node& conv_node = builder.AddQLinearConvNode(input_arg, .01f, 135, + conv1_weight_arg, .02f, 126, + conv_output_arg, .37f, 131); + conv_node.AddAttribute("pads", std::vector{1, 1, 1, 1}); + + Node& split_node = builder.AddNode("Split", {conv_output_arg}, {split_output1_arg, split_output2_arg}); + split_node.AddAttribute("axis", static_cast(axis)); + + Node& qlconcat_node = builder.AddQLinearConcatLike( + "QLinearConcat", qlconcat_output_arg, .37f, 131, + {{split_output1_arg, .37f, uint8_t(131)}, {split_output2_arg, .37f, uint8_t(131)}}); + qlconcat_node.AddAttribute("axis", static_cast(axis)); + + auto* conv2_weight_arg = NhwcMakeInitializer(builder, {17, conv1_output_channels, 3, 3}); + builder.AddQLinearConvNode(qlconcat_output_arg, .43f, 126, + conv2_weight_arg, .02f, 126, + output_arg, .37f, 131); + }; + + auto check_nhwc_graph = [&](InferenceSessionWrapper& session) { + auto op_to_count = CountOpsInGraph(session.GetGraph()); + EXPECT_EQ(op_to_count["com.microsoft.QLinearConv"], 2); + EXPECT_EQ(op_to_count["Transpose"], 2); + }; + + TransformerTester(build_test_case, + check_nhwc_graph, + TransformerLevel::Level2, + TransformerLevel::Level3); + } +} + TEST(NhwcTransformerTests, ConvPad) { std::vector pad_modes{"constant", "reflect", "edge"}; for (const auto& mode : pad_modes) {