Re-implement caffe old-style extractors with extractor extensions (#3675)

* move crop extractor

* Add concat_ext.py

* Add roipooling_ext.py

* Add roipooling_ext

* Add scale extractor

* Add scale extractor

* Add bn_ext.py and dropout_ext.py

* Add bn_ext.py and dropout_ext.py

* Add bn_ext.py and dropout_ext.py

* Fix bn.ext.py

* Sort fix

* Fix bn_test.py

* rename to batchnorm_ext

* Add bn_ext

* Fix batchnorm_ext.py

* small fix

* Small fix

model-optimizer/extensions/front/caffe/batchnorm_ext.py

@@ -0,0 +1,53 @@
+"""
+ Copyright (C) 2018-2021 Intel Corporation
+
+ Licensed under the Apache License, Version 2.0 (the "License");
+ you may not use this file except in compliance with the License.
+ You may obtain a copy of the License at
+
+      http://www.apache.org/licenses/LICENSE-2.0
+
+ Unless required by applicable law or agreed to in writing, software
+ distributed under the License is distributed on an "AS IS" BASIS,
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ See the License for the specific language governing permissions and
+ limitations under the License.
+"""
+import numpy as np
+
+from extensions.ops.BatchNormInference import BatchNormInference
+from mo.front.caffe.extractors.utils import embed_input
+from mo.front.extractor import FrontExtractorOp
+
+
+class BatchNormalizationExtractor(FrontExtractorOp):
+    op = 'batchnorm'
+    enabled = True
+
+    @classmethod
+    def extract(cls, node):
+        eps = node.pb.batch_norm_param.eps
+        attrs = {
+           'eps': eps
+        }
+        pb_model = None if not node.soft_get('model_pb', None) else node.model_pb
+        if pb_model:
+            blobs = pb_model.blobs
+            assert len(blobs) >= 2, 'BatchNorm accepts not less then two input blobs'
+            mean = np.array(blobs[0].data)
+            variance = np.array(blobs[1].data)
+
+            if len(blobs) == 3:
+                scale = blobs[2].data[0]
+                if scale != 0:
+                    scale = 1.0 / scale
+                mean *= scale
+                variance *= scale
+
+            embed_input(attrs, 1, 'gamma', np.ones(mean.shape), 'gamma')
+            embed_input(attrs, 2, 'beta', np.zeros(variance.shape), 'beta')
+            embed_input(attrs, 3, 'mean', mean, 'biases')
+            embed_input(attrs, 4, 'variance', variance, 'weights')
+
+        BatchNormInference.update_node_stat(node, attrs)
+        return cls.enabled

model-optimizer/extensions/front/caffe/bn.py

@@ -1,5 +1,5 @@
 """
- Copyright (C) 2018-2020 Intel Corporation
+ Copyright (C) 2018-2021 Intel Corporation
 
  Licensed under the Apache License, Version 2.0 (the "License");
  you may not use this file except in compliance with the License.
@@ -27,7 +27,7 @@ class BNToScaleShift(FrontReplacementOp):
     """
     Replaces BN layer with ScaleShift.
     """
-    op = "batchNormInference"
+    op = "BN"
     enabled = True
 
     def replace_op(self, graph: Graph, node: Node):
@@ -35,6 +35,7 @@ class BNToScaleShift(FrontReplacementOp):
 
         param = graph.node[node.id]['pb'].bn_param
         pb_model = graph.node[node.id]['model_pb']
+
         blobs = pb_model.blobs
 
         if len(blobs) != 4:

model-optimizer/extensions/front/caffe/bn_ext.py

@@ -0,0 +1,28 @@
+"""
+ Copyright (C) 2018-2021 Intel Corporation
+
+ Licensed under the Apache License, Version 2.0 (the "License");
+ you may not use this file except in compliance with the License.
+ You may obtain a copy of the License at
+
+      http://www.apache.org/licenses/LICENSE-2.0
+
+ Unless required by applicable law or agreed to in writing, software
+ distributed under the License is distributed on an "AS IS" BASIS,
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ See the License for the specific language governing permissions and
+ limitations under the License.
+"""
+
+from extensions.ops.BN import BN
+from mo.front.extractor import FrontExtractorOp
+
+
+class BNExtractor(FrontExtractorOp):
+    op = 'BN'
+    enabled = True
+
+    @classmethod
+    def extract(cls, node):
+        BN.update_node_stat(node, {})
+        return cls.enabled

model-optimizer/extensions/front/caffe/bn_test.py

@@ -1,5 +1,5 @@
 """
- Copyright (C) 2018-2020 Intel Corporation
+ Copyright (C) 2018-2021 Intel Corporation
 
  Licensed under the Apache License, Version 2.0 (the "License");
  you may not use this file except in compliance with the License.
@@ -13,9 +13,8 @@
  See the License for the specific language governing permissions and
  limitations under the License.
 """
-import unittest
-
 import numpy as np
+import unittest
 
 from extensions.front.caffe.bn import BNToScaleShift
 from mo.graph.graph import Node
@@ -47,7 +46,7 @@ class TestBNReplacer(unittest.TestCase):
                                  FakeParam('data', shift)])
         nodes = [
             ('input', {'kind': 'op', 'type': 'Identity', 'op': 'Identity'}),
-            ('bn', {'type': None, 'kind': 'op', 'op': 'batchNormInference', 'pb': bn_pb, 'model_pb': bn_bin}),
+            ('bn', {'type': None, 'kind': 'op', 'op': 'BN', 'pb': bn_pb, 'model_pb': bn_bin}),
             ('output', {'kind': 'op', 'type': 'Identity', 'op': 'Identity'}),
         ]
         edges = [

model-optimizer/extensions/front/caffe/concat_ext.py

@@ -0,0 +1,32 @@
+"""
+ Copyright (C) 2018-2021 Intel Corporation
+
+ Licensed under the Apache License, Version 2.0 (the "License");
+ you may not use this file except in compliance with the License.
+ You may obtain a copy of the License at
+
+      http://www.apache.org/licenses/LICENSE-2.0
+
+ Unless required by applicable law or agreed to in writing, software
+ distributed under the License is distributed on an "AS IS" BASIS,
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ See the License for the specific language governing permissions and
+ limitations under the License.
+"""
+
+from mo.front.extractor import FrontExtractorOp
+from mo.ops.concat import Concat
+
+
+class ConcatFrontExtractor(FrontExtractorOp):
+    op = 'concat'
+    enabled = True
+
+    @classmethod
+    def extract(cls, node):
+        pb = node.pb
+        mapping_rule = {
+           'axis': pb.concat_param.axis,
+        }
+        Concat.update_node_stat(node, mapping_rule)
+        return cls.enabled

model-optimizer/extensions/front/caffe/crop_ext.py

@@ -0,0 +1,39 @@
+"""
+ Copyright (C) 2018-2021 Intel Corporation
+
+ Licensed under the Apache License, Version 2.0 (the "License");
+ you may not use this file except in compliance with the License.
+ You may obtain a copy of the License at
+
+      http://www.apache.org/licenses/LICENSE-2.0
+
+ Unless required by applicable law or agreed to in writing, software
+ distributed under the License is distributed on an "AS IS" BASIS,
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ See the License for the specific language governing permissions and
+ limitations under the License.
+"""
+
+from mo.front.common.partial_infer.crop import crop_infer
+from mo.front.extractor import FrontExtractorOp
+from mo.ops.crop import Crop
+
+
+class CropFrontExtractor(FrontExtractorOp):
+    op = 'Crop'
+    enabled = True
+
+    @classmethod
+    def extract(cls, node):
+        proto_layer = node.pb
+        param = proto_layer.crop_param
+        mapping_rule = {
+            'type': 'Crop',
+            'axis': param.axis,
+            'offset': param.offset,
+            'dim': None,  # set in infer
+            'infer': crop_infer
+        }
+        # update the attributes of the node
+        Crop.update_node_stat(node, mapping_rule)
+        return cls.enabled

model-optimizer/extensions/front/caffe/crop_ext_test.py

@@ -0,0 +1,66 @@
+"""
+ Copyright (C) 2018-2021 Intel Corporation
+
+ Licensed under the Apache License, Version 2.0 (the "License");
+ you may not use this file except in compliance with the License.
+ You may obtain a copy of the License at
+
+      http://www.apache.org/licenses/LICENSE-2.0
+
+ Unless required by applicable law or agreed to in writing, software
+ distributed under the License is distributed on an "AS IS" BASIS,
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ See the License for the specific language governing permissions and
+ limitations under the License.
+"""
+
+import unittest
+from unittest.mock import patch
+
+from extensions.front.caffe.crop_ext import CropFrontExtractor
+from mo.front.common.partial_infer.crop import crop_infer
+from mo.ops.crop import Crop
+from mo.ops.op import Op
+from mo.utils.unittest.extractors import FakeMultiParam
+from mo.utils.unittest.graph import FakeNode
+
+
+class FakeCropProtoLayer:
+    def __init__(self, val):
+        self.crop_param = val
+
+
+class TestCropExt(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        Op.registered_ops['Crop'] = Crop
+
+    def test_da_no_pb_no_ml(self):
+        self.assertRaises(AttributeError, CropFrontExtractor.extract, None)
+
+    @patch('mo.front.caffe.collect_attributes')
+    def test_crop_ext(self, collect_attributes_mock):
+        params = {
+            'axis': 0,
+            'offset': 0,
+        }
+        collect_attributes_mock.return_value = {
+            **params,
+            'test': 54,
+            'test2': 'test3'
+        }
+        fake_pl = FakeCropProtoLayer(FakeMultiParam(params))
+        fake_node = FakeNode(fake_pl, None)
+
+        CropFrontExtractor.extract(fake_node)
+
+        exp_res = {
+            'type': 'Crop',
+            'axis': 0,
+            'offset': 0,
+            'dim': None,  # set in infer
+            'infer': crop_infer
+        }
+
+        for key in exp_res.keys():
+            self.assertEqual(exp_res[key], fake_node[key])

model-optimizer/extensions/front/caffe/dropout_ext.py

@@ -0,0 +1,29 @@
+"""
+ Copyright (C) 2018-2021 Intel Corporation
+
+ Licensed under the Apache License, Version 2.0 (the "License");
+ you may not use this file except in compliance with the License.
+ You may obtain a copy of the License at
+
+      http://www.apache.org/licenses/LICENSE-2.0
+
+ Unless required by applicable law or agreed to in writing, software
+ distributed under the License is distributed on an "AS IS" BASIS,
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ See the License for the specific language governing permissions and
+ limitations under the License.
+"""
+
+from extensions.ops.identity import Identity
+from mo.front.extractor import FrontExtractorOp
+from mo.graph.graph import Node
+
+
+class DropoutFrontExtractor(FrontExtractorOp):
+    op = 'dropout'
+    enabled = True
+
+    @classmethod
+    def extract(cls, node: Node):
+        Identity.update_node_stat(node, {})
+        return cls.enabled

model-optimizer/extensions/front/caffe/roipooling_ext.py

@@ -0,0 +1,35 @@
+"""
+ Copyright (C) 2018-2021 Intel Corporation
+
+ Licensed under the Apache License, Version 2.0 (the "License");
+ you may not use this file except in compliance with the License.
+ You may obtain a copy of the License at
+
+      http://www.apache.org/licenses/LICENSE-2.0
+
+ Unless required by applicable law or agreed to in writing, software
+ distributed under the License is distributed on an "AS IS" BASIS,
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ See the License for the specific language governing permissions and
+ limitations under the License.
+"""
+
+from mo.front.extractor import FrontExtractorOp
+from mo.ops.roipooling import ROIPooling
+
+
+class ROIPoolingFrontExtractor(FrontExtractorOp):
+    op = 'roipooling'
+    enabled = True
+
+    @classmethod
+    def extract(cls, node):
+        param = node.pb.roi_pooling_param
+        attrs =  {
+            'pooled_h': param.pooled_h,
+            'pooled_w': param.pooled_w,
+            'spatial_scale': param.spatial_scale,
+        }
+        
+        ROIPooling.update_node_stat(node, attrs)
+        return cls.enabled

model-optimizer/extensions/front/caffe/scale_ext.py

@@ -0,0 +1,55 @@
+"""
+ Copyright (C) 2018-2021 Intel Corporation
+
+ Licensed under the Apache License, Version 2.0 (the "License");
+ you may not use this file except in compliance with the License.
+ You may obtain a copy of the License at
+
+      http://www.apache.org/licenses/LICENSE-2.0
+
+ Unless required by applicable law or agreed to in writing, software
+ distributed under the License is distributed on an "AS IS" BASIS,
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ See the License for the specific language governing permissions and
+ limitations under the License.
+"""
+
+import numpy as np
+
+from mo.front.caffe.extractors.utils import embed_input, weights_biases
+from mo.front.common.partial_infer.elemental import copy_shape_infer
+from mo.front.extractor import FrontExtractorOp
+from mo.ops.scale_shift import ScaleShiftOp
+from mo.utils.utils import NamedAttrsClass
+
+
+class ScaleFrontExtractor(FrontExtractorOp):
+    op = 'scale'
+    enabled = True
+
+    @classmethod
+    def extract(cls, node):
+        pb = node.pb
+        model = node.model_pb
+        param = pb.scale_param
+        attrs = {
+            'axis': param.axis,
+        }
+        
+        if model is None and len(pb.bottom) == 1:
+            # default weights and biases for scale layer if the caffemodel file doesn't contain them
+            model = NamedAttrsClass({'blobs': np.array([NamedAttrsClass({'data': np.array([1])}),
+                                                 NamedAttrsClass({'data': np.array([0])})])})
+        # scale with 1 input and 1 or 2 blobs
+        if model and len(model.blobs) != 0 and len(pb.bottom) == 1:
+            attrs.update(weights_biases(param.bias_term, model))
+        # 2 inputs + bias
+        elif len(pb.bottom) == 2 and param.bias_term:
+            if model is None or len(model.blobs) == 0:
+                # default bias for scale layer with 2 inputs if the caffemodel file doesn't contain them
+                model = NamedAttrsClass({'blobs': np.array([NamedAttrsClass({'data': np.array([0])})])})
+
+            embed_input(attrs, 1, 'biases', model.blobs[0].data)
+        ScaleShiftOp.update_node_stat(node, attrs)
+        return cls.enabled
+

model-optimizer/extensions/ops/BN.py

@@ -0,0 +1,35 @@
+"""
+ Copyright (C) 2018-2021 Intel Corporation
+
+ Licensed under the Apache License, Version 2.0 (the "License");
+ you may not use this file except in compliance with the License.
+ You may obtain a copy of the License at
+
+      http://www.apache.org/licenses/LICENSE-2.0
+
+ Unless required by applicable law or agreed to in writing, software
+ distributed under the License is distributed on an "AS IS" BASIS,
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ See the License for the specific language governing permissions and
+ limitations under the License.
+"""
+
+from mo.graph.graph import Graph
+from mo.ops.op import Op
+
+
+class BN(Op):
+    """
+    BN operation comes from caffe and will be replaced by BNToScaleShift FrontReplacer.
+    """
+    op = 'BN'
+    enabled = False
+
+    def __init__(self, graph: Graph, attrs: dict):
+        super().__init__(graph, {
+            'type': None,
+            'op': self.op,
+            'in_ports_count': 5,
+            'out_ports_count': 1,
+            'infer': None
+        }, attrs)