PyTorch MO extension (openvinotoolkit#10)

* PyTorch MO extension

* Add tests

* Test RetinaNet

* Use master OpenVINO

* Move PyTorch tests to Azure

.ci/azure/linux.yml

@@ -79,12 +79,23 @@ jobs:
       wget https://github.com/ninja-build/ninja/releases/download/v1.10.0/ninja-linux.zip
       unzip ninja-linux.zip
       sudo cp -v ninja /usr/local/bin/
+      python3 -m pip install --upgrade pip
+      python3 -m pip install cython
     workingDirectory: $(WORK_DIR)
     displayName: 'Install dependencies'
 
   - task: CMake@1
     inputs:
-      cmakeArgs: -GNinja -DVERBOSE_BUILD=ON -DCMAKE_BUILD_TYPE=$(BUILD_TYPE) -DIE_EXTRA_MODULES=$(REPO_DIR)/modules $(OPENVINO_REPO_DIR)
+      cmakeArgs: >
+        -GNinja
+        -DVERBOSE_BUILD=ON
+        -DCMAKE_BUILD_TYPE=$(BUILD_TYPE)
+        -DIE_EXTRA_MODULES=$(REPO_DIR)/modules
+        -DENABLE_PYTHON=ON
+        -DPYTHON_EXECUTABLE=/usr/bin/python3
+        -DPYTHON_LIBRARY=/usr/lib/x86_64-linux-gnu/libpython3.6m.so
+        -DPYTHON_INCLUDE_DIR=/usr/include/python3.6
+        $(OPENVINO_REPO_DIR)
       workingDirectory: $(BUILD_DIR)
 
   - script: ninja
@@ -99,3 +110,16 @@ jobs:
     env:
       LD_LIBRARY_PATH: $(BIN_DIR)/lib
     displayName: 'Java tests'
+
+  - script: |
+      sudo apt-get install -y python3-setuptools
+      python3 -m pip install -r requirements.txt
+    workingDirectory: $(REPO_DIR)/modules/mo_pytorch/test
+    displayName: 'Install PyTorch tests dependencies'
+
+  - script: python3 run_tests.py
+    env:
+      LD_LIBRARY_PATH: $(BIN_DIR)/lib
+      PYTHONPATH: "$(OPENVINO_REPO_DIR)/model-optimizer:$(BIN_DIR)/lib/python_api/python3.6:$(REPO_DIR)/modules/mo_pytorch"
+    workingDirectory: $(REPO_DIR)/modules/mo_pytorch/test
+    displayName: 'PyTorch conversion tests'

modules/README.md

@@ -7,3 +7,4 @@ $ cmake -DIE_EXTRA_MODULES=<openvino_contrib>/modules -DBUILD_<module name>=OFF
 ```
 
 * [**java_api**](./java_api): Inference Engine Java API -- provides Java wrappers for Inference Engine public API.
+* [**mo_pytorch**](./mo_pytorch): PyTorch extensions for Model Optimizer -- native PyTorch to OpenVINO IR converter

modules/mo_pytorch/.gitignore

@@ -0,0 +1 @@
+__pycache__

modules/mo_pytorch/README.md

@@ -0,0 +1,38 @@
+# PyTorch extension for Model Optimizer
+
+This module let you convert PyTorch models directly to OpenVINO IR without using ONNX
+
+## Usage
+
+1. Clone repository
+
+    ```bash
+    git clone --depth 1 https://github.com/openvinotoolkit/openvino_contrib
+    ```
+
+2. Setup environment
+
+    ```bash
+    source /opt/intel/openvino_<VERSION>/bin/setupvars.sh
+    export PYTHONPATH=openvino_contrib/modules/mo_pytorch:$PYTHONPATH
+    ```
+
+3. Convert PyTorch model to OpenVINO IR
+
+    ```python
+    import torchvision.models as models
+
+    # Create model
+    model = models.alexnet(pretrained=True)
+
+    # Convert to OpenVINO IR
+    import mo_pytorch
+    mo_pytorch.convert(model, input_shape=[1, 3, 227, 227], model_name='alexnet')
+    ```
+
+## Supported networks
+
+* `torchvision.models.alexnet`
+* `torchvision.models.resnet18`
+* `torchvision.models.segmentation.deeplabv3_resnet50`
+* `Detectron2 RetinaNet`

modules/mo_pytorch/mo_extensions/front/pytorch/activation_ext.py

@@ -0,0 +1,36 @@
+"""
+ Copyright (C) 2018-2020 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.activation_ops import *
+from mo.front.extractor import FrontExtractorOp
+
+class ReLUExtractor(FrontExtractorOp):
+    op = 'ReLU'
+    enabled = True
+
+    @classmethod
+    def extract(cls, node):
+        ReLU.update_node_stat(node)
+        return cls.enabled
+
+
+class SigmoidExtractor(FrontExtractorOp):
+    op = 'Sigmoid'
+    enabled = True
+
+    @classmethod
+    def extract(cls, node):
+        Sigmoid.update_node_stat(node)
+        return cls.enabled

modules/mo_pytorch/mo_extensions/front/pytorch/batchnorm.py

@@ -0,0 +1,45 @@
+"""
+ Copyright (C) 2018-2020 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.elemental import copy_shape_infer
+from mo.front.common.replacement import FrontReplacementOp
+from mo.graph.graph import Graph, Node
+from mo.ops.const import Const
+from mo.ops.op import Op
+
+
+class BatchNorm(Op):
+    op = 'BatchNorm'
+
+    def __init__(self, graph: Graph, attrs: dict):
+        mandatory_props = {
+            'type': self.op,
+            'op': self.op,
+            'eps': None,
+            'infer': copy_shape_infer,
+            'in_ports_count': 5,
+            'out_ports_count': 1,
+        }
+        super().__init__(graph, mandatory_props, attrs)
+
+class BatchNorm2d(FrontReplacementOp):
+    op = 'BatchNorm2d'
+    enabled = True
+
+    def replace_op(self, graph: Graph, node: Node):
+        inputs = [node.in_node(i) for i in range(5)]
+        bn = BatchNorm(graph, dict(name=node.name, eps=node.module.eps)).create_node(inputs)
+        return [bn.id]

modules/mo_pytorch/mo_extensions/front/pytorch/const_ext.py

@@ -0,0 +1,32 @@
+"""
+ Copyright (C) 2018-2020 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.const import Const
+
+
+class ConstExtractor(FrontExtractorOp):
+    op = 'Const'
+    enabled = True
+
+    @classmethod
+    def extract(cls, node):
+        value = node.value
+        attrs = {
+            'data_type': value.dtype,
+            'value': value
+        }
+        Const.update_node_stat(node, attrs)
+        return cls.enabled

modules/mo_pytorch/mo_extensions/front/pytorch/conv_ext.py

@@ -0,0 +1,64 @@
+"""
+ Copyright (C) 2018-2020 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.common.partial_infer.utils import int64_array
+from mo.front.extractor import FrontExtractorOp
+from mo.ops.convolution import Convolution
+from mo.utils.error import Error
+
+
+class Conv2dFrontExtractor(FrontExtractorOp):
+    op = 'Conv2d'
+    enabled = True
+
+    @classmethod
+    def extract(cls, node):
+        # Extract pads attribute
+        pads = np.array(node.module.padding, dtype=np.int64).reshape(1, 2)
+        pads = np.repeat(pads, 2, axis=0)
+        final_pads = np.array([[0, 0], [0, 0], *pads], dtype=np.int64)
+
+        # Extract strides attribute
+        strides = node.module.stride
+        final_strides = np.array([1, 1, *strides], dtype=np.int64)
+
+        # Extract strides attribute
+        dilations = node.module.dilation
+        if isinstance(dilations, int):
+            dilations = [dilations, dilations]
+        final_dilations = np.array([1, 1, *dilations], dtype=np.int64)
+
+        attrs = {
+            'op': 'Conv2D',  # Note: should be 2D but not 2d
+            'pad': final_pads,
+            'stride': final_strides,
+            'dilation': final_dilations,
+            'group': 1,
+            'kernel_spatial': np.array(node.module.kernel_size, dtype=np.int64),
+
+            'input_feature_channel': 1,
+            'output_feature_channel': 0,
+
+            'channel_dims': np.array([1], dtype=np.int64),
+            'batch_dims': np.array([0], dtype=np.int64),
+            'layout': 'NCHW',
+        }
+
+        # update the attributes of the node
+        Convolution.update_node_stat(node, attrs)
+        return cls.enabled

modules/mo_pytorch/mo_extensions/front/pytorch/detection_output_ext.py

@@ -0,0 +1,40 @@
+"""
+ Copyright (C) 2018-2020 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.common.partial_infer.utils import int64_array
+from mo.front.extractor import FrontExtractorOp
+from extensions.ops.DetectionOutput import DetectionOutput
+from mo.utils.error import Error
+
+
+class DetectionOutputExtractor(FrontExtractorOp):
+    op = 'DetectionOutput'
+    enabled = True
+
+    @classmethod
+    def extract(cls, node):
+        attrs = {
+            'variance_encoded_in_target': int(node.module.variance_encoded_in_target),
+            'nms_threshold': node.module.nms_threshold,
+            'confidence_threshold': node.module.confidence_threshold,
+            'top_k': node.module.top_k,
+            'keep_top_k': node.module.keep_top_k,
+            'code_type': node.module.code_type,
+        }
+        DetectionOutput.update_node_stat(node, attrs)
+        return cls.enabled

modules/mo_pytorch/mo_extensions/front/pytorch/dropout_ext.py

@@ -0,0 +1,29 @@
+"""
+ Copyright (C) 2018-2020 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.utils.error import Error
+
+
+class DropoutFrontExtractor(FrontExtractorOp):
+    op = 'Dropout'
+    enabled = True
+
+    @classmethod
+    def extract(cls, node):
+        Identity.update_node_stat(node)
+        return cls.enabled

modules/mo_pytorch/mo_extensions/front/pytorch/elementwise_ext.py

@@ -0,0 +1,52 @@
+"""
+ Copyright (C) 2018-2020 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.elementwise import *
+from mo.front.extractor import FrontExtractorOp
+from mo.graph.graph import Node
+from mo.ops.eltwise_n import EltwiseNAdd, EltwiseNMax
+from mo.ops.power import AttributedPower
+
+
+class AddFrontExtractor(FrontExtractorOp):
+    op = 'Add'
+    enabled = True
+
+    @classmethod
+    def extract(cls, node: Node):
+        Add.update_node_stat(node)
+        return cls.enabled
+
+
+class SubFrontExtractor(FrontExtractorOp):
+    op = 'Sub'
+    enabled = True
+
+    @classmethod
+    def extract(cls, node: Node):
+        Sub.update_node_stat(node)
+        return cls.enabled
+
+
+class MulFrontExtractor(FrontExtractorOp):
+    op = 'Mul'
+    enabled = True
+
+    @classmethod
+    def extract(cls, node: Node):
+        Mul.update_node_stat(node)
+        return cls.enabled