From 08113b80cc3c4bd0c446cd479d9a919f2a23d00d Mon Sep 17 00:00:00 2001 From: Tracy Sharpe <42477615+tracysh@users.noreply.github.com> Date: Mon, 20 Jan 2020 16:35:03 -0800 Subject: [PATCH] Optimize BatchNormalization to NCHWc Conv (#2855) Update the NCHWc transformer to convert BatchNormalization ops to NCHWc convolutions when the input tensor is already in NCHWc. --- .../core/optimizer/nchwc_transformer.cc | 131 ++++++++++++++++-- .../test/optimizer/nchwc_optimizer_test.cc | 61 +++++++- 2 files changed, 180 insertions(+), 12 deletions(-) diff --git a/onnxruntime/core/optimizer/nchwc_transformer.cc b/onnxruntime/core/optimizer/nchwc_transformer.cc index a75debc761..5d108a0c3d 100644 --- a/onnxruntime/core/optimizer/nchwc_transformer.cc +++ b/onnxruntime/core/optimizer/nchwc_transformer.cc @@ -118,6 +118,7 @@ class NchwcTransformerImpl { void TransformAdd(Node& node); void TransformConcat(Node& node); void TransformActivation(Node& node); + void TransformBatchNormalization(Node& node); Graph& graph_; @@ -310,7 +311,6 @@ void NchwcTransformerImpl::TransformConv(Node& node) { } const size_t nchwc_block_size = MlasNchwcGetBlockSize(); - const int64_t nchwc_output_channels = (output_channels + nchwc_block_size - 1) & ~(nchwc_block_size - 1); bool do_reorder_input = true; @@ -386,9 +386,7 @@ void NchwcTransformerImpl::TransformConv(Node& node) { nchwc_conv_W_tensor_proto.add_dims(conv_W->dims()[i]); } - graph_.AddInitializedTensor(nchwc_conv_W_tensor_proto); - - nchwc_conv_W_arg = &graph_.GetOrCreateNodeArg(nchwc_conv_W_tensor_proto.name(), nullptr); + nchwc_conv_W_arg = &graph_utils::AddInitializer(graph_, nchwc_conv_W_tensor_proto); filters_map->emplace(input_defs[1], nchwc_conv_W_arg); } @@ -409,13 +407,11 @@ void NchwcTransformerImpl::TransformConv(Node& node) { nchwc_conv_B_tensor_proto.set_data_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT); nchwc_conv_B_tensor_proto.set_name(graph_.GenerateNodeArgName("reorder")); - nchwc_conv_B_tensor_proto.set_raw_data(aligned_bias.data(), aligned_bias.size() * sizeof(float)); + nchwc_conv_B_tensor_proto.set_raw_data(aligned_bias.data(), nchwc_output_channels * sizeof(float)); nchwc_conv_B_tensor_proto.add_dims(nchwc_output_channels); - graph_.AddInitializedTensor(nchwc_conv_B_tensor_proto); - - nchwc_conv_B_arg = &graph_.GetOrCreateNodeArg(nchwc_conv_B_tensor_proto.name(), nullptr); + nchwc_conv_B_arg = &graph_utils::AddInitializer(graph_, nchwc_conv_B_tensor_proto); aligned_biases_.emplace(input_defs[2], nchwc_conv_B_arg); } } @@ -666,6 +662,121 @@ void NchwcTransformerImpl::TransformActivation(Node& node) { } } +// Transform BatchNormalization to a depthwise separable 1x1 convolution. This +// enables reuse of the existing NCHWc convolution operator and other fusions +// such as BatchNormalization+Relu using Conv+Relu. +void NchwcTransformerImpl::TransformBatchNormalization(Node& node) { + auto& input_defs = node.MutableInputDefs(); + auto& output_defs = node.MutableOutputDefs(); + + // Don't transform the node if the input is not already in NCHWc format. + auto it = nchwc_args_.find(input_defs[0]); + if (it == nchwc_args_.end()) { + return; + } + auto* nchwc_input = it->second.get(); + + // Require that BatchNormalization-7 uses spatial normalization. + const auto* spatial_attr = graph_utils::GetNodeAttribute(node, "spatial"); + if (spatial_attr != nullptr && utils::HasInt(*spatial_attr) && spatial_attr->i() != 1) { + return; + } + + const auto* epsilon_attr = graph_utils::GetNodeAttribute(node, "epsilon"); + if (epsilon_attr == nullptr || !utils::HasFloat(*epsilon_attr)) { + return; + } + float epsilon = static_cast(epsilon_attr->f()); + + const int64_t channels = nchwc_input->channels_; + + auto get_bn_tensor_proto = [this, channels](const std::string& input_name) { + const auto* tensor_proto = graph_utils::GetConstantInitializer(graph_, input_name); + if (tensor_proto != nullptr) { + if ((tensor_proto->data_type() != ONNX_NAMESPACE::TensorProto_DataType_FLOAT) || + (tensor_proto->dims_size() != 1) || + (tensor_proto->dims(0) != channels)) { + tensor_proto = nullptr; + } + } + return tensor_proto; + }; + + const auto* bn_scale_tensor_proto = get_bn_tensor_proto(input_defs[1]->Name()); + if (bn_scale_tensor_proto == nullptr) { + return; + } + const auto* bn_B_tensor_proto = get_bn_tensor_proto(input_defs[2]->Name()); + if (bn_B_tensor_proto == nullptr) { + return; + } + const auto* bn_mean_tensor_proto = get_bn_tensor_proto(input_defs[3]->Name()); + if (bn_mean_tensor_proto == nullptr) { + return; + } + const auto* bn_var_tensor_proto = get_bn_tensor_proto(input_defs[4]->Name()); + if (bn_var_tensor_proto == nullptr) { + return; + } + + auto bn_scale = onnxruntime::make_unique(*bn_scale_tensor_proto); + auto bn_B = onnxruntime::make_unique(*bn_B_tensor_proto); + auto bn_mean = onnxruntime::make_unique(*bn_mean_tensor_proto); + auto bn_var = onnxruntime::make_unique(*bn_var_tensor_proto); + + // Calculate the scale and bias for the replacement convolution. + bn_var->add(epsilon); + bn_var->sqrt(); + bn_scale->div(*bn_var); + bn_mean->mul(*bn_scale); + bn_B->sub(*bn_mean); + + const size_t nchwc_block_size = MlasNchwcGetBlockSize(); + const int64_t nchwc_channels = (channels + nchwc_block_size - 1) & ~(nchwc_block_size - 1); + + std::vector padded_buffer(nchwc_channels); + + std::copy_n(bn_scale->data(), channels, padded_buffer.data()); + + ONNX_NAMESPACE::TensorProto nchwc_conv_W_tensor_proto; + nchwc_conv_W_tensor_proto.set_data_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT); + nchwc_conv_W_tensor_proto.set_name(graph_.GenerateNodeArgName("bn_scale")); + nchwc_conv_W_tensor_proto.set_raw_data(padded_buffer.data(), nchwc_channels * sizeof(float)); + nchwc_conv_W_tensor_proto.add_dims(nchwc_channels); + nchwc_conv_W_tensor_proto.add_dims(1); + nchwc_conv_W_tensor_proto.add_dims(1); + nchwc_conv_W_tensor_proto.add_dims(1); + + auto* nchwc_conv_W_arg = &graph_utils::AddInitializer(graph_, nchwc_conv_W_tensor_proto); + + std::copy_n(bn_B->data(), channels, padded_buffer.data()); + + ONNX_NAMESPACE::TensorProto nchwc_conv_B_tensor_proto; + nchwc_conv_B_tensor_proto.set_data_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT); + nchwc_conv_B_tensor_proto.set_name(graph_.GenerateNodeArgName("bn_B")); + nchwc_conv_B_tensor_proto.set_raw_data(padded_buffer.data(), nchwc_channels * sizeof(float)); + nchwc_conv_B_tensor_proto.add_dims(nchwc_channels); + + auto* nchwc_conv_B_arg = &graph_utils::AddInitializer(graph_, nchwc_conv_B_tensor_proto); + + // Create the replacement node. + std::string nchwc_node_name = graph_.GenerateNodeName(output_defs[0]->Name() + "_bn_nchwc"); + Node& nchwc_node = graph_.AddNode(nchwc_node_name, + "Conv", + nchwc_node_name, + {nchwc_input->nchwc_arg_, nchwc_conv_W_arg, nchwc_conv_B_arg}, + output_defs, + nullptr, + kMSNchwcDomain); + nchwc_node.SetExecutionProviderType(node.GetExecutionProviderType()); + nchwc_node.AddAttribute("group", nchwc_channels); + + nchwc_input->remaining_original_uses_--; + + CreateNchwcArgument(node, nchwc_node, channels, nchwc_input->shape_); + removed_nodes_.push_front(node.Index()); +} + void NchwcTransformerImpl::Transform(Node& node) { if (graph_utils::IsSupportedOptypeVersionAndDomain(node, "Conv", {1, 11}) || graph_utils::IsSupportedOptypeVersionAndDomain(node, "FusedConv", {1}, kMSDomain)) { @@ -688,6 +799,8 @@ void NchwcTransformerImpl::Transform(Node& node) { TransformConcat(node); } else if (graph_utils::IsSupportedOptypeVersionAndDomain(node, "Relu", {6})) { TransformActivation(node); + } else if (graph_utils::IsSupportedOptypeVersionAndDomain(node, "BatchNormalization", {7, 9})) { + TransformBatchNormalization(node); } } @@ -713,8 +826,8 @@ void NchwcTransformerImpl::Finalize(bool& modified) { {output_original_arg}, nullptr, kMSNchwcDomain); - reorder_output_node.AddAttribute("channels", nchwc_output.second->channels_); reorder_output_node.SetExecutionProviderType(kCpuExecutionProvider); + reorder_output_node.AddAttribute("channels", nchwc_output.second->channels_); } } diff --git a/onnxruntime/test/optimizer/nchwc_optimizer_test.cc b/onnxruntime/test/optimizer/nchwc_optimizer_test.cc index 3eea1bf0dd..3a8bb81904 100644 --- a/onnxruntime/test/optimizer/nchwc_optimizer_test.cc +++ b/onnxruntime/test/optimizer/nchwc_optimizer_test.cc @@ -39,7 +39,7 @@ class NchwcInferenceSession : public InferenceSession { }; struct NchwcTestHelper { - NchwcTestHelper(Graph& graph) : graph_(graph), fill_value_(0) { + NchwcTestHelper(Graph& graph) : graph_(graph), fill_value_(0), per_sample_tolerance_(0.0) { } NodeArg* MakeInput(const std::vector& shape, const ONNX_NAMESPACE::TypeProto& type_proto) { @@ -157,6 +157,7 @@ struct NchwcTestHelper { NameMLValMap feeds_; std::vector output_names_; int fill_value_; + double per_sample_tolerance_; }; void NchwcOptimizerTester(const std::function& build_test_case, @@ -210,12 +211,11 @@ void NchwcOptimizerTester(const std::function& bu ASSERT_TRUE(num_outputs == level3_fetches.size()); for (size_t i = 0; i < num_outputs; i++) { - double per_sample_tolerance = 0.0; double relative_per_sample_tolerance = 0.0; std::pair ret = CompareOrtValue(level3_fetches[i], level2_fetches[i], - per_sample_tolerance, + helper.per_sample_tolerance_, relative_per_sample_tolerance, false); EXPECT_EQ(ret.first, COMPARE_RESULT::SUCCESS); @@ -932,6 +932,61 @@ TEST(NchwcOptimizerTests, TensorAlignment) { NchwcOptimizerTester(build_test_case, check_nchwc_graph); } +TEST(NchwcOptimizerTests, BatchNormalization) { + auto build_test_case = [&](NchwcTestHelper& helper) { + auto* input_arg = helper.MakeInput({1, 1, 23, 21}); + auto* conv1_output_arg = helper.MakeIntermediate(); + auto* conv2_output_arg = helper.MakeIntermediate(); + auto* output_arg = helper.MakeOutput(); + + // Using a channel count not aligned to the block size to verify handling + // of unaligned data. + helper.AddConvNode(input_arg, conv1_output_arg, {34, 1, 3, 3}); + helper.AddConvNode(input_arg, conv2_output_arg, {34, 1, 3, 3}); + + auto* add_output_arg = helper.MakeIntermediate(); + helper.AddNode("Add", {conv1_output_arg, conv2_output_arg}, {add_output_arg}); + + std::vector bn_scale(34); + std::vector bn_bias(34); + std::vector bn_mean(34); + std::vector bn_var(34); + + for (int i = 0; i < 34; i++) { + bn_scale[i] = static_cast((i % 5) + 1) * 0.01f; + bn_bias[i] = static_cast(i - 17) * 0.25f; + bn_mean[i] = static_cast(i % 7) * 0.001f; + bn_var[i] = static_cast((i % 9) + 1) * 0.001f; + } + + auto* bn_scale_arg = helper.MakeInitializer({34}, bn_scale); + auto* bn_bias_arg = helper.MakeInitializer({34}, bn_bias); + auto* bn_mean_arg = helper.MakeInitializer({34}, bn_mean); + auto* bn_var_arg = helper.MakeInitializer({34}, bn_var); + + auto* bn_output_arg = helper.MakeIntermediate(); + helper.AddNode("BatchNormalization", {add_output_arg, bn_scale_arg, bn_bias_arg, bn_mean_arg, bn_var_arg}, {bn_output_arg}); + helper.AddNode("Relu", {bn_output_arg}, {output_arg}); + + // Override the sample tolerance for this test. By default, the NCHWc + // tests generate bit identical results when run with and without + // optimizations, but the BatchNormalizationtransform does introduce + // small bit differences. + helper.per_sample_tolerance_ = .000125; + }; + + auto check_nchwc_graph = [&](NchwcInferenceSession& session) { + auto op_to_count = session.CountOpsInGraph(); + EXPECT_EQ(op_to_count["nchwc.Conv"], 3); + EXPECT_EQ(op_to_count["nchwc.ReorderInput"], 0); + EXPECT_EQ(op_to_count["nchwc.ReorderOutput"], 1); + EXPECT_EQ(op_to_count["BatchNormalization"], 0); + EXPECT_EQ(op_to_count["Relu"], 0); + }; + + NchwcOptimizerTester(build_test_case, check_nchwc_graph); +} + #endif } // namespace test