diff --git a/model-optimizer/mo/ops/convolution.py b/model-optimizer/mo/ops/convolution.py index 31b469719f2450..2c331169254a4a 100644 --- a/model-optimizer/mo/ops/convolution.py +++ b/model-optimizer/mo/ops/convolution.py @@ -20,7 +20,6 @@ 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 @@ def calc_convolution(input_spatial_shape, stride_spatial_shape, pad_spatial_shap 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 in incorrect input_shape") + float_spatial_val_wo_stride = float_array(spatial_val_wo_stride) return float_spatial_val_wo_stride / stride_spatial_shape + 1 diff --git a/model-optimizer/mo/ops/convolution_test.py b/model-optimizer/mo/ops/convolution_test.py index 0c612f805c0e8a..a9573cbc238a63 100644 --- a/model-optimizer/mo/ops/convolution_test.py +++ b/model-optimizer/mo/ops/convolution_test.py @@ -21,6 +21,7 @@ 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 @@ def test_conv_infer_3D_convolution(self): 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 95ab11b85e4b78..a2c0ffdfd21f37 100644 --- a/model-optimizer/mo/ops/pooling.py +++ b/model-optimizer/mo/ops/pooling.py @@ -17,12 +17,10 @@ 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.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) @@ -96,9 +94,14 @@ def infer(node: Node): 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 in incorrect input_shape") + + output_spatial_shape = np.array(rounding(np.array(padded_spatial_shape, dtype=np.float) / stride_spatial), + dtype=np.int64) + 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 556859ce1098f1..456e31b92669b8 100644 --- a/model-optimizer/mo/ops/pooling_test.py +++ b/model-optimizer/mo/ops/pooling_test.py @@ -21,6 +21,7 @@ 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 @@ def test_pooling_infer_no_shape(self): 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)