added check to avoid IR generation in case of wrong input shape (#2127)

* added check to avoid IR generation in case of wrong input shape

* review changes

diff --git a/model-optimizer/mo/ops/convolution.py b/model-optimizer/mo/ops/convolution.py
index 31b4697..95471b0 100644 (file)
--- a/model-optimizer/mo/ops/convolution.py
+++ b/model-optimizer/mo/ops/convolution.py
@@ -20,7 +20,6 @@ import numpy as np
 
 from mo.front.common.partial_infer.utils import int64_array, float_array, mark_input_bins, assign_dims_to_weights, \
     tf_window_op_pad_infer
-from mo.front.extractor import spatial_getter
 from mo.front.onnx.extractors.utils import get_backend_pad
 from mo.graph.graph import Node, Graph
 from mo.ops.op import Op, PermuteAttrs
@@ -72,6 +71,11 @@ class Convolution(Op):
             Verified to be applicable for both Caffe and ONNX.
         '''
         spatial_val_wo_stride = input_spatial_shape + pad_spatial_shape - kernel_extent
+
+        if np.any(spatial_val_wo_stride < 0):
+            raise Error("Data after padding has dimension less than window size. " +
+                        "Possible reason of error is incorrectly specified model input shape(s).")
+
         float_spatial_val_wo_stride = float_array(spatial_val_wo_stride)
         return float_spatial_val_wo_stride / stride_spatial_shape + 1
 
@@ -118,8 +122,8 @@ class Convolution(Op):
                 log.error('Cannot reshape kernel due to not all required attrs was set to {} node'.format(node.id))
                 return
             # layout for Convolution weights is OIHW
-            kernel_shape = np.array([node.output, input_shape[node.channel_dims].item() / node.group,
-                                    *[node.kernel_spatial[i] for i in range(len(node.kernel_spatial))]], dtype=np.int64)
+            kernel_shape = int64_array([node.output, input_shape[node.channel_dims].item() / node.group,
+                                       *[node.kernel_spatial[i] for i in range(len(node.kernel_spatial))]])
             if node.type == 'Deconvolution':  # layout for Deconvolution weights is IOHW
                 kernel_shape[[0, 1]] = kernel_shape[[1, 0]]
                 #node.input_feature_channel, node.output_feature_channel = node.output_feature_channel, node.input_feature_channel
@@ -163,7 +167,7 @@ class Convolution(Op):
         if not node.has_valid('stride'):
             node['stride'] = np.full([len(input_shape)], 1, dtype=np.int64)
         if not node.has_valid('pad'):
-            node['pad'] = np.array([[0, 0]] * len(input_shape), dtype=np.int64)
+            node['pad'] = int64_array([[0, 0]] * len(input_shape))
         node['pad_spatial_shape'] = node.pad[node.spatial_dims]
 
         if not node.has_valid('output_padding'):

diff --git a/model-optimizer/mo/ops/convolution_test.py b/model-optimizer/mo/ops/convolution_test.py
index 0c612f8..a9573cb 100644 (file)
--- a/model-optimizer/mo/ops/convolution_test.py
+++ b/model-optimizer/mo/ops/convolution_test.py
@@ -21,6 +21,7 @@ import numpy as np
 from mo.front.common.partial_infer.utils import int64_array
 from mo.graph.graph import Node
 from mo.ops.convolution import Convolution
+from mo.utils.error import Error
 from mo.utils.unittest.extractors import FakeValue
 from mo.utils.unittest.graph import build_graph
 
@@ -378,3 +379,31 @@ class TestConvolutionPartialInfer(unittest.TestCase):
         self.assertTrue(np.array_equal(int64_array([[0, 0], [0, 0], [0, 0]]), conv_node.pad_spatial_shape))
         # Check resulting output shape
         self.assertTrue(np.array_equal(int64_array([1, 64, 16, 218, 218]), conv_output.shape))
+
+    def test_caffe_conv2d_infer_wrong_input_shape(self):
+        graph = build_graph(nodes_attributes,
+                            [('conv_input', 'conv_node'),
+                             ('conv_weights', 'conv_node'),
+                             ('conv_node', 'conv_output'),
+                             ('conv_output', 'op_output')
+                             ],
+                            {'conv_output': {'shape': None},
+                             'conv_input': {'shape': np.array([1, 3, 1, 1])},
+                             'conv_weights': {'shape': np.array([64, 3, 3, 3]),
+                                              'dim_attrs': ['spatial_dims', 'channel_dims', 'batch_dims', 'axis']},
+                             'conv_node': {'pad_spatial_shape': np.array([[0, 0], [0, 0]]),
+                                           'conv_pad': np.array([[0, 0], [0, 0], [0, 0], [0, 0]]),
+                                           'dilation': np.array([1, 1, 1, 1]), 'bias_addable': True, 'bias_term': False,
+                                           'output_spatial_shape': None, 'output_shape': None,
+                                           'stride': np.array([1, 1, 1, 1]), 'group': 1,
+                                           'kernel_spatial_idx': np.array([2, 3]),
+                                           'input_feature_channel': 1,
+                                           'output_feature_channel': 0,
+                                           'output': 64, 'kernel_spatial': np.array([3, 3]),
+                                           'spatial_dims': np.array([2, 3]), 'channel_dims': np.array([1]),
+                                           'batch_dims': np.array([0])}
+                             })
+
+        conv_node = Node(graph, 'conv_node')
+        with self.assertRaises(Error):
+            Convolution.infer(conv_node)

diff --git a/model-optimizer/mo/ops/pooling.py b/model-optimizer/mo/ops/pooling.py
index 95ab11b..0d65a92 100644 (file)
--- a/model-optimizer/mo/ops/pooling.py
+++ b/model-optimizer/mo/ops/pooling.py
@@ -16,13 +16,11 @@
 
 import numpy as np
 
-from mo.front.common.partial_infer.utils import tf_window_op_pad_infer
-from mo.front.extractor import attr_getter
-# from mo.front.common.partial_infer.pooling import pool_explicit_padding_infer
-from mo.front.extractor import spatial_getter
+from mo.front.common.partial_infer.utils import tf_window_op_pad_infer, int64_array, float_array
 from mo.front.onnx.extractors.utils import get_backend_pad
 from mo.graph.graph import Node, Graph
 from mo.ops.op import Op, PermuteAttrs
+from mo.utils.error import Error
 
 
 class Pooling(Op):
@@ -30,10 +28,10 @@ class Pooling(Op):
 
     def __init__(self, graph: Graph, attrs: dict):
         super().__init__(graph, {
-            'type': __class__.op,
-            'op': __class__.op,
+            'type': self.op,
+            'op': self.op,
             'version': 'opset1',
-            'infer': __class__.infer,
+            'infer': self.infer,
             'in_ports_count': 1,
             'out_ports_count': 1,
         }, attrs)
@@ -68,11 +66,11 @@ class Pooling(Op):
 
         # Setting default pad and stride attrs in case of None specified
         if not node.has_valid('pad'):
-            node['pad'] = np.array([[0, 0] for x in range(len(input_shape))], dtype=np.int64)
+            node['pad'] = int64_array([[0, 0] for x in range(len(input_shape))])
         if not node.has_valid('pad_spatial_shape'):
             node['pad_spatial_shape'] = node.pad[node.spatial_dims]
         if not node.has_valid('stride'):
-            node['stride'] = np.array([1 for x in range(len(input_shape))], dtype=np.int64)
+            node['stride'] = int64_array([1 for x in range(len(input_shape))])
 
         if node.has_and_set('global_pool'):
             node['window'] = np.zeros(len(input_shape), dtype=np.int64)
@@ -96,9 +94,13 @@ class Pooling(Op):
             rounding = np.floor
             if node.soft_get('pooling_convention') == 'full' or node.soft_get('rounding_type') == 'ceil':
                 rounding = np.ceil
-            output_spatial_shape = np.array(rounding(
-                np.array(input_spatial_shape + pad_spatial_shape - window_spatial_shape,
-                         dtype=np.float) / stride_spatial), dtype=np.int64) + 1
+
+            padded_spatial_shape = input_spatial_shape + pad_spatial_shape - window_spatial_shape
+            if np.any(padded_spatial_shape < 0):
+                raise Error("Data after padding has dimension less than window size. " +
+                            "Possible reason of error is incorrectly specified model input shape(s).")
+
+            output_spatial_shape = int64_array(rounding(float_array(padded_spatial_shape) / stride_spatial)) + 1
 
             original_pads = np.array([i[1] for i in node.pad_spatial_shape])
 

diff --git a/model-optimizer/mo/ops/pooling_test.py b/model-optimizer/mo/ops/pooling_test.py
index 556859c..456e31b 100644 (file)
--- a/model-optimizer/mo/ops/pooling_test.py
+++ b/model-optimizer/mo/ops/pooling_test.py
@@ -21,6 +21,7 @@ import numpy as np
 from mo.graph.graph import Node
 from mo.ops.pooling import Pooling
 from mo.utils.unittest.graph import build_graph
+from mo.utils.error import Error
 
 nodes_attributes = {'node_1': {'value': None, 'kind': 'data'},
                     'pool': {'type': 'Pooling', 'value': None, 'kind': 'op'},
@@ -129,3 +130,26 @@ class TestPoolingPartialInfer(unittest.TestCase):
         Pooling.infer(pool_node)
         res_shape = graph.node['node_2']['shape']
         self.assertIsNone(res_shape)
+
+    def test_pooling_infer_wrong_input_shape(self):
+        graph = build_graph(nodes_attributes,
+                            [('node_1', 'pool'),
+                             ('pool', 'node_2'),
+                             ('node_2', 'op_output')
+                             ],
+                            {'node_2': {'shape': None},
+                             'node_1': {'shape': np.array([1, 3, 1, 1])},
+                             'pool': {'window': np.array([1, 1, 5, 5]), 'stride': np.array([1, 1, 2, 2]),
+                                      'pad': np.array([[0, 0], [0, 0], [1, 1], [1, 1]]),
+                                      'pad_spatial_shape': np.array([[1, 1], [1, 1]]),
+                                      'pool_method': 'avg', 'exclude_pad': 'false', 'global_pool': 0,
+                                      'output_spatial_shape': None, 'output_shape': None,
+                                      'kernel_spatial': np.array([3, 3]), 'spatial_dims': np.array([2, 3]),
+                                      'channel_dims': np.array([1]), 'batch_dims': np.array([0]),
+                                      'pooling_convention': 'full'}
+                             })
+
+        pool_node = Node(graph, 'pool')
+
+        with self.assertRaises(Error):
+            Pooling.infer(pool_node)