mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-30 20:18:08 +00:00
Fix NCHWc BatchNormalization regression (#2903)
Fix the BatchNormalization optimization in the NCHWC optimizer. If the node has the optional training outputs specified, then skip the transform.
This commit is contained in:
parent
ba336b5583
commit
17b72d5578
2 changed files with 70 additions and 43 deletions
|
|
@ -669,6 +669,11 @@ void NchwcTransformerImpl::TransformBatchNormalization(Node& node) {
|
|||
auto& input_defs = node.MutableInputDefs();
|
||||
auto& output_defs = node.MutableOutputDefs();
|
||||
|
||||
// Bail out if the node has the optional training outputs specified.
|
||||
if (output_defs.size() > 1) {
|
||||
return;
|
||||
}
|
||||
|
||||
// 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()) {
|
||||
|
|
|
|||
|
|
@ -933,58 +933,80 @@ TEST(NchwcOptimizerTests, TensorAlignment) {
|
|||
}
|
||||
|
||||
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();
|
||||
auto test_case = [&](bool training_outputs) {
|
||||
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});
|
||||
// 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});
|
||||
auto* add_output_arg = helper.MakeIntermediate();
|
||||
helper.AddNode("Add", {conv1_output_arg, conv2_output_arg}, {add_output_arg});
|
||||
|
||||
std::vector<float> bn_scale(34);
|
||||
std::vector<float> bn_bias(34);
|
||||
std::vector<float> bn_mean(34);
|
||||
std::vector<float> bn_var(34);
|
||||
std::vector<float> bn_scale(34);
|
||||
std::vector<float> bn_bias(34);
|
||||
std::vector<float> bn_mean(34);
|
||||
std::vector<float> bn_var(34);
|
||||
|
||||
for (int i = 0; i < 34; i++) {
|
||||
bn_scale[i] = static_cast<float>((i % 5) + 1) * 0.01f;
|
||||
bn_bias[i] = static_cast<float>(i - 17) * 0.25f;
|
||||
bn_mean[i] = static_cast<float>(i % 7) * 0.001f;
|
||||
bn_var[i] = static_cast<float>((i % 9) + 1) * 0.001f;
|
||||
}
|
||||
for (int i = 0; i < 34; i++) {
|
||||
bn_scale[i] = static_cast<float>((i % 5) + 1) * 0.01f;
|
||||
bn_bias[i] = static_cast<float>(i - 17) * 0.25f;
|
||||
bn_mean[i] = static_cast<float>(i % 7) * 0.001f;
|
||||
bn_var[i] = static_cast<float>((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_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});
|
||||
auto* bn_output_arg = helper.MakeIntermediate();
|
||||
std::vector<NodeArg*> bn_output_args = {bn_output_arg};
|
||||
if (training_outputs) {
|
||||
bn_output_args.push_back(helper.MakeIntermediate());
|
||||
bn_output_args.push_back(helper.MakeIntermediate());
|
||||
bn_output_args.push_back(helper.MakeIntermediate());
|
||||
bn_output_args.push_back(helper.MakeIntermediate());
|
||||
}
|
||||
helper.AddNode("BatchNormalization", {add_output_arg, bn_scale_arg, bn_bias_arg, bn_mean_arg, bn_var_arg}, bn_output_args);
|
||||
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;
|
||||
// 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();
|
||||
if (training_outputs) {
|
||||
EXPECT_EQ(op_to_count["nchwc.Conv"], 2);
|
||||
EXPECT_EQ(op_to_count["BatchNormalization"], 1);
|
||||
EXPECT_EQ(op_to_count["Relu"], 1);
|
||||
} else {
|
||||
EXPECT_EQ(op_to_count["nchwc.Conv"], 3);
|
||||
EXPECT_EQ(op_to_count["BatchNormalization"], 0);
|
||||
EXPECT_EQ(op_to_count["Relu"], 0);
|
||||
}
|
||||
EXPECT_EQ(op_to_count["nchwc.ReorderInput"], 0);
|
||||
EXPECT_EQ(op_to_count["nchwc.ReorderOutput"], 1);
|
||||
};
|
||||
|
||||
NchwcOptimizerTester(build_test_case, check_nchwc_graph);
|
||||
};
|
||||
|
||||
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);
|
||||
// Verify that a batch normalization node can be converted to a convolution
|
||||
// if the input tensor is already in NCHWc format. However, this transform
|
||||
// should be skipped if the batch normalization node has the optional training
|
||||
// outputs supplied.
|
||||
test_case(false);
|
||||
test_case(true);
|
||||
}
|
||||
|
||||
#endif
|
||||
|
|
|
|||
Loading…
Reference in a new issue