MO now shows a hint when --disable_nhwc_to_nchw is needed (#6350)

* MO now shows a hint when --disable_nhwc_to_nchw is needed

* Updated message to be printed using log.error

* Small code refactoring

model-optimizer/extensions/load/tf/loader.py

@@ -22,7 +22,7 @@ from mo.front.common.register_custom_ops import update_extractors_with_extension
 from mo.front.extractor import restore_edges, extract_node_attrs, remove_control_dependency_inputs, add_outputs_identity
 from mo.front.tf.extractor import get_tf_edges, create_tf_edge, tf_op_extractor, tf_op_extractors
 from mo.front.tf.loader import load_tf_graph_def, protobuf2nx
-from mo.graph.graph import Graph
+from mo.graph.graph import Graph, Node
 from mo.utils import tensorboard_util
 from mo.utils.error import Error
 from mo.utils.telemetry_utils import send_op_names_info, send_shapes_info, send_framework_info
@@ -42,12 +42,12 @@ class TFLoader(Loader):
                 tf_v1.load_op_library(library)
 
         graph_def, variables_values, framework = load_tf_graph_def(graph_file_name=argv.input_model,
-                                                        is_binary=not argv.input_model_is_text,
-                                                        checkpoint=argv.input_checkpoint,
-                                                        user_output_node_names_list=argv.output,
-                                                        model_dir=argv.saved_model_dir,
-                                                        meta_graph_file=argv.input_meta_graph,
-                                                        saved_model_tags=argv.saved_model_tags)
+                                                                   is_binary=not argv.input_model_is_text,
+                                                                   checkpoint=argv.input_checkpoint,
+                                                                   user_output_node_names_list=argv.output,
+                                                                   model_dir=argv.saved_model_dir,
+                                                                   meta_graph_file=argv.input_meta_graph,
+                                                                   saved_model_tags=argv.saved_model_tags)
         send_framework_info(framework)
 
         try:
@@ -100,5 +100,46 @@ class TFLoader(Loader):
 
         graph.check_empty_graph('protobuf2nx. It may happen due to problems with loaded model')
         extract_node_attrs(graph, lambda node: tf_op_extractor(node, check_for_duplicates(tf_op_extractors)))
+
+        # try to detect layout from the nodes of the graph. If there are no convolution nodes in N(D)HWC layout then we
+        # consider that the graph is in NCHW layout and no layout conversion should be performed
+        if not argv.disable_nhwc_to_nchw and not argv.silent and not graph_or_sub_graph_has_nhwc_ops(graph):
+            log.error('The TensorFlow model does not contain Convolution operations with N(D)HWC layout. Most likely '
+                      'the model should be converted using additional "--disable_nhwc_to_nchw" command line parameter '
+                      'which disables model layout conversion inside the Model Optimizer.', extra={'is_warning': True})
+
         send_op_names_info(framework, graph)
         send_shapes_info(framework, graph)
+
+
+def is_node_layout_nhwc(node: Node):
+    """
+    Check the layout attribute of specific operations and return True if any of them has layout NHWC.
+    :param node: Node to check
+    :return: Boolean result of the check
+    """
+    if node.soft_get('op') in ["Conv2D", "DepthwiseConv2dNative", "Conv3D", "Conv2DBackpropInput",
+                               "Conv3DBackpropInputV2"]:
+        if node.soft_get('layout') in ["NHWC", "NDHWC"]:
+            log.debug('Detected convolution node with NHWC layout: "{}"'.format(node.soft_get('name', node.id)))
+            return True
+    return False
+
+
+def graph_or_sub_graph_has_nhwc_ops(graph: Graph):
+    """
+    Checks that a graph or any sub-graph (inside Loop) operation contains nodes with NHWC layout.
+    :param graph: main graph to check
+    :return: Boolean result of the check
+    """
+    NHWC_conv_detected = False
+    for node in graph.get_op_nodes():
+        if is_node_layout_nhwc(node):
+            NHWC_conv_detected = True
+            break
+
+        # for the Loop node we need to check that the body does not contain marker ops as well
+        if node.op == 'Loop':
+            NHWC_conv_detected |= graph_or_sub_graph_has_nhwc_ops(node.body)
+            # TODO check for If op when it is implemented
+    return NHWC_conv_detected

model-optimizer/unit_tests/extensions/load/tf/loader_test.py

@@ -0,0 +1,68 @@
+# Copyright (C) 2018-2021 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+import unittest
+
+import numpy as np
+
+from extensions.load.tf.loader import graph_or_sub_graph_has_nhwc_ops
+from unit_tests.utils.graph import build_graph, result, regular_op, const, connect_front
+
+
+class TFLoaderTest(unittest.TestCase):
+    @staticmethod
+    def build_conv_graph():
+        nodes = {
+            **const('weights', np.random.randn(1, 1, 1, 1)),
+            **regular_op('input', {'op': 'Parameter'}),
+            **regular_op('conv', {'op': 'Conv2D', 'layout': 'NHWC'}),
+            **result('result'),
+        }
+        edges = [*connect_front('input', '0:conv'),
+                 *connect_front('weights', '1:conv'),
+                 *connect_front('conv:0', 'result'),
+                 ]
+        graph = build_graph(nodes, edges)
+
+        graph.stage = 'front'
+        return graph
+
+    @staticmethod
+    def build_parameter_result_graph():
+        nodes = {
+            **regular_op('input', {'op': 'Parameter'}),
+            **result('result'),
+        }
+        edges = [*connect_front('input', '0:result'),
+                 ]
+        graph = build_graph(nodes, edges)
+        graph.stage = 'front'
+        return graph
+
+    @staticmethod
+    def build_loop_graph(body_graph):
+        # create fake Loop operation
+        nodes = {
+            **regular_op('input', {'op': 'Parameter'}),
+            **regular_op('loop', {'op': 'Loop', 'body': body_graph}),
+            **result('result'),
+        }
+        edges = [*connect_front('input', '0:loop'),
+                 *connect_front('loop:0', 'result'),
+                 ]
+        graph = build_graph(nodes, edges)
+        graph.stage = 'front'
+        return graph
+
+    def test_convolution_main_graph(self):
+        self.assertTrue(graph_or_sub_graph_has_nhwc_ops(self.build_conv_graph()))
+
+    def test_convolution_loop_body_graph(self):
+        self.assertTrue(graph_or_sub_graph_has_nhwc_ops(self.build_loop_graph(self.build_conv_graph())))
+
+    def test_no_convolution_main_graph(self):
+        self.assertFalse(graph_or_sub_graph_has_nhwc_ops(self.build_parameter_result_graph()))
+
+    def test_no_convolution_main_and_sub_graph(self):
+        self.assertFalse(graph_or_sub_graph_has_nhwc_ops(self.build_loop_graph(self.build_parameter_result_graph())))
+