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:
Tracy Sharpe 2020-01-23 18:54:11 -08:00 committed by GitHub
parent ba336b5583
commit 17b72d5578
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 70 additions and 43 deletions

View file

@ -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()) {

View file

@ -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