Add support for ONNX BatchNorm-7 and -9 (#5465)

openvinotoolkit · Apr 30, 2021 · bcb67bf · bcb67bf
1 parent 03ca3d1
commit bcb67bf
Show file tree

Hide file tree

Showing 3 changed files with 36 additions and 4 deletions.
diff --git a/ngraph/frontend/onnx_import/src/op/batch_norm.cpp b/ngraph/frontend/onnx_import/src/op/batch_norm.cpp
@@ -18,6 +18,7 @@ namespace ngraph
         {
             namespace set_1
             {
+                // This version supports ONNX BatchNormalization-1 and BatchNormalization-6
                 OutputVector batch_norm(const Node& node)
                 {
                     OutputVector inputs{node.get_ng_inputs()};
@@ -27,11 +28,10 @@ namespace ngraph
                     Output<ngraph::Node> mean;
                     Output<ngraph::Node> var;
 
-                    std::int64_t is_test{node.get_attribute_value<std::int64_t>("is_test", 1)};
                     double epsilon{node.get_attribute_value<double>("epsilon", 1e-5)};
 
-                    // TODO: Implement learning mode support
-                    // float momentum{node.get_attribute_value<float>("momentum", 0.9f)};
+                    // Currently only BatchNormalization inference mode is supported by OpenVINO
+                    std::int64_t is_test{node.get_attribute_value<std::int64_t>("is_test", 1)};
                     CHECK_VALID_NODE(node, is_test, "only 'is_test' mode is supported.");
 
                     // optional outputs
@@ -55,9 +55,34 @@ namespace ngraph
                     throw ngraph_error(
                         "Cannot create nGraph batch norm with unsupported number of inputs");
                 }
-
             } // namespace set_1
 
+            namespace set_7
+            {
+                // This version supports ONNX BatchNormalization-7 and BatchNormalization-9
+                OutputVector batch_norm(const Node& node)
+                {
+                    OutputVector inputs{node.get_ng_inputs()};
+                    auto x = inputs.at(0);
+                    auto scale = inputs.at(1);
+                    auto bias = inputs.at(2);
+                    auto mean = inputs.at(3);
+                    auto var = inputs.at(4);
+
+                    double epsilon{node.get_attribute_value<double>("epsilon", 1e-5)};
+                    // Attribute "spatial" is ignored, as we only support inference mode of
+                    // BatchNormalization
+
+                    CHECK_VALID_NODE(node,
+                                     node.get_outputs_size() == 1,
+                                     "Training mode of BatchNormalization is not supported.");
+
+                    return {std::make_shared<default_opset::BatchNormInference>(
+                        x, scale, bias, mean, var, epsilon)};
+                }
+
+            } // namespace set_7
+
         } // namespace op
 
     } // namespace onnx_import

diff --git a/ngraph/frontend/onnx_import/src/op/batch_norm.hpp b/ngraph/frontend/onnx_import/src/op/batch_norm.hpp
@@ -19,6 +19,12 @@ namespace ngraph
 
             } // namespace set_1
 
+            namespace set_7
+            {
+                OutputVector batch_norm(const Node& node);
+
+            } // namespace set_7
+
         } // namespace op
 
     } // namespace onnx_import

diff --git a/ngraph/frontend/onnx_import/src/ops_bridge.cpp b/ngraph/frontend/onnx_import/src/ops_bridge.cpp
@@ -316,6 +316,7 @@ namespace ngraph
             REGISTER_OPERATOR("Atanh", 1, atanh);
             REGISTER_OPERATOR("AveragePool", 1, average_pool);
             REGISTER_OPERATOR("BatchNormalization", 1, batch_norm);
+            REGISTER_OPERATOR("BatchNormalization", 7, batch_norm);
             REGISTER_OPERATOR("BitShift", 1, bitshift);
             REGISTER_OPERATOR("Cast", 1, cast);
             REGISTER_OPERATOR("Ceil", 1, ceil);