OpenVINO extensions: custom operations support (#481)

Author: likholat

.ci/azure/linux.yml

@@ -49,6 +49,7 @@ jobs:
     BIN_DIR: $(OPENVINO_REPO_DIR)/bin/intel64/$(BUILD_TYPE)
     INSTALL_DIR: $(WORK_DIR)/install_pkg
     SETUPVARS: $(INSTALL_DIR)/setupvars.sh
+    CUSTOM_OP_LIB: $(BIN_DIR)/libuser_ov_extensions.so
     GRADLE_VER: 7.1.1
 
   steps:
@@ -93,7 +94,7 @@ jobs:
       set -e
       sudo -E $(OPENVINO_REPO_DIR)/install_build_dependencies.sh
       # Move jdk into contrib install_build_dependencies.sh
-      sudo apt --assume-yes install default-jdk
+      sudo apt --assume-yes install default-jdk libopencv-dev
       # Install gradle for Java
       sudo apt -y --no-install-recommends install unzip
       wget https://services.gradle.org/distributions/gradle-$(GRADLE_VER)-bin.zip
@@ -155,3 +156,20 @@ jobs:
       done
     workingDirectory: $(REPO_DIR)/modules/java_api
     displayName: 'Java tests'
+
+  - script: |
+      python3 -m pip install --user virtualenv
+      python3 -m virtualenv -p /usr/bin/python3.8 .env3
+      source .env3/bin/activate
+      python -m pip install --upgrade pip
+      python -m pip install -r $(REPO_DIR)/modules/custom_operations/tests/requirements.txt
+      cd ${OPENVINO_REPO_DIR}/tools && python -m pip install mo/
+    workingDirectory: $(WORK_DIR)
+    displayName: 'Create custom ops env'
+
+  - script: |
+      . $(SETUPVARS)
+      source  $(WORK_DIR)/.env3/bin/activate
+      python -m pytest tests/run_tests.py
+    workingDirectory: $(REPO_DIR)/modules/custom_operations
+    displayName: 'Custom ops tests'

.ci/azure/windows.yml

@@ -45,11 +45,12 @@ jobs:
     MODELS_PATH: $(REPO_DIR)\..\testdata
     WORK_DIR: $(Pipeline.Workspace)\_w
     BUILD_DIR: D:\build
-    BIN_DIR: $(OPENVINO_REPO_DIR)\bin\intel64
+    BIN_DIR: $(OPENVINO_REPO_DIR)\bin\intel64\$(BUILD_TYPE)
     MSVS_VARS_PATH: C:\Program Files (x86)\Microsoft Visual Studio\2019\Enterprise\VC\Auxiliary\Build\vcvars64.bat
     MSVC_COMPILER_PATH: C:\Program Files (x86)\Microsoft Visual Studio\2019\Enterprise\VC\Tools\MSVC\14.24.28314\bin\Hostx64\x64\cl.exe
     INSTALL_DIR: $(WORK_DIR)\install_pkg
     SETUPVARS: $(INSTALL_DIR)\setupvars.bat
+    CUSTOM_OP_LIB: $(BIN_DIR)\user_ov_extensions.dll
     GRADLE_VER: 7.1.1
 
   steps:
@@ -96,6 +97,8 @@ jobs:
       powershell -command "Expand-Archive -Force ninja-win.zip"
       powershell -command "Invoke-WebRequest https://services.gradle.org/distributions/gradle-$(GRADLE_VER)-bin.zip -OutFile gradle-$(GRADLE_VER)-bin.zip"
       powershell -command "Expand-Archive -Force gradle-$(GRADLE_VER)-bin.zip"
+      python -m pip install --upgrade pip
+      python -m pip install -r $(OPENVINO_REPO_DIR)\src\bindings\python\src\compatibility\openvino\requirements-dev.txt
     workingDirectory: $(WORK_DIR)
     displayName: 'Install dependencies'
 
@@ -107,6 +110,7 @@ jobs:
         -DENABLE_INTEL_MYRIAD_COMMON=OFF ^
         -DENABLE_CPPLINT=OFF ^
         -DOPENVINO_EXTRA_MODULES=$(REPO_DIR)/modules ^
+        -DENABLE_PYTHON=ON ^
         -DCMAKE_C_COMPILER:PATH="$(MSVC_COMPILER_PATH)" ^
         -DCMAKE_CXX_COMPILER:PATH="$(MSVC_COMPILER_PATH)" ^
         $(OPENVINO_REPO_DIR)

README.md

@@ -12,6 +12,7 @@ This list gives an overview of all modules available inside the contrib reposito
 * [**nvidia_plugin**](./modules/nvidia_plugin): NVIDIA GPU Plugin -- allows to perform deep neural networks inference on NVIDIA GPUs using CUDA, using OpenVINO API.
 * [**java_api**](./modules/java_api): Inference Engine Java API -- provides Java wrappers for Inference Engine public API.
 * [**Azure Video Analyzer**](./modules/ovms_ai_extension/): Azure Video Analyzer Extension -- enables exchange of video frames and inference results between [Azure Video Analyzer (AVA)](https://docs.microsoft.com/en-us/azure/azure-video-analyzer/video-analyzer-docs/overview) and OpenVINO™ Model Server.
+* [**custom_operations**](./modules/custom_operations/): Collection of Custom Operations -- implement Custom Operations with OpenVINO Extensibility Mechanism.
 
 ## How to build OpenVINO with extra modules
 You can build OpenVINO, so it will include the modules from this repository. Contrib modules are under constant development and it is recommended to use them alongside the master branch or latest releases of OpenVINO.
@@ -35,6 +36,7 @@ Additional build instructions are available for the following modules:
 * [**arm_plugin**](./modules/arm_plugin/README.md)
 * [**nvidia_plugin**](./modules/nvidia_plugin/README.md)
 * [**ovms_ai_extension**](./modules/ovms_ai_extension/README.md)
+* [**custom_operations**](./modules/custom_operations/README.md)
 
 ## Update the repository documentation
 In order to keep a clean overview containing all contributed modules, the following files need to be created/adapted:

modules/custom_operations/CMakeLists.txt

@@ -0,0 +1,8 @@
+# Copyright (C) 2022 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+#
+
+cmake_minimum_required(VERSION 3.13)
+project(openvino_extensions)
+
+add_subdirectory(user_ie_extensions)

modules/custom_operations/README.md

@@ -0,0 +1,66 @@
+# Collection of Custom Operations using OpenVINO Extensibility Mechanism
+
+This module provides a guide and the implementation of a few custom operations in the Intel OpenVINO runtime using its [Extensibility Mechanism](https://docs.openvino.ai/latest/openvino_docs_Extensibility_UG_Intro.html).
+
+There are some use cases when OpenVINO Custom Operations could be applicable:
+
+* There is an ONNX model which contains an operation not supported by OpenVINO.
+* You have a PyTorch model, which could be converted to ONNX, with an operation not supported by OpenVINO.
+* You want to replace a subgraph for ONNX model with one custom operation which would be supported by OpenVINO.
+
+More specifically, here we implement custom OpenVINO operations that add support for the following native PyTorch operation:
+
+* [torch.fft](examples/fft)
+
+And other custom operations introduced by third-party frameworks:
+
+* [calculate_grid](/examples/calculate_grid) and [sparse_conv](/examples/sparse_conv) from [Open3D](https://github.com/isl-org/Open3D)
+* [complex_mul](/examples/complex_mul) from [DIRECT](https://github.com/NKI-AI/direct)
+
+You can find more information about how to create and use OpenVINO Extensions to facilitate mapping of custom operations from framework model representation to OpenVINO representation [here](https://docs.openvino.ai/latest/openvino_docs_Extensibility_UG_Frontend_Extensions.html).
+
+
+## Build custom OpenVINO operation extension library
+
+The C++ code implementing the custom operation is in the `user_ie_extensions` directory. You'll have to build an "extension library" from this code so that it can be loaded at runtime. The steps below describe the build process:
+
+1. Install [OpenVINO Runtime for C++](https://docs.openvino.ai/latest/openvino_docs_install_guides_install_dev_tools.html#for-c-developers).
+
+2. Build the library:
+
+```bash
+cd user_ie_extensions
+mkdir build && cd build
+cmake .. -DCMAKE_BUILD_TYPE=Release && cmake --build . --parallel 4
+```
+
+If you need to build only some operations specify them with the `-DCUSTOM_OPERATIONS` option:
+```bash
+cmake .. -DCMAKE_BUILD_TYPE=Release -DCUSTOM_OPERATIONS=complex_mul\;fft
+```
+
+- Please note that [OpenCV](https://opencv.org/) installation is required to build an extension for the [fft](examples/fft) operation. Other extentions still can be built without OpenCV.
+
+You also could build the extension library [while building OpenVINO](../../README.md).
+
+## Load and use custom OpenVINO operation extension library
+
+You can use the custom OpenVINO operations implementation by loading it into the OpenVINO `Core` object at runtime. Then, load the model from the ONNX file with the `read_model()` API. Here's how to do that in Python:
+
+```python
+from openvino.runtime import Core
+
+# Create Core and register user extension
+core = Core()
+core.add_extension('/path/to/libuser_ov_extensions.so')
+
+# Load model from .onnx file directly
+model = core.read_model('model.onnx')
+compiled_model = core.compile_model(model, 'CPU')
+```
+
+You also can get OpenVINO IR model with Model Optimizer, just use extra `--extension` flag to specify a path to custom extensions:
+
+```bash
+mo --input_model model.onnx --extension /path/to/libuser_ov_extensions.so
+```

modules/custom_operations/examples/calculate_grid/calculate_grid.py

@@ -0,0 +1,27 @@
+# Copyright (C) 2018-2022 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class CalculateGrid(torch.autograd.Function):
+    @staticmethod
+    def symbolic(g, in_positions):
+        return g.op("CalculateGrid", in_positions)
+
+    @staticmethod
+    def forward(self, in_positions):
+        filter = torch.Tensor([[-1, -1, -1], [-1, -1, 0], [-1, 0, -1], [-1, 0, 0],
+                               [0, -1, -1], [0, -1, 0], [0, 0, -1],
+                               [0, 0, 0]]).to(in_positions.device)
+
+        out_pos = in_positions.long().repeat(1, filter.shape[0]).reshape(-1, 3)
+        filter = filter.repeat(in_positions.shape[0], 1)
+
+        out_pos = out_pos + filter
+        out_pos = out_pos[out_pos.min(1).values >= 0]
+        out_pos = out_pos[(~((out_pos.long() % 2).bool()).any(1))]
+        out_pos = torch.unique(out_pos, dim=0)
+
+        return out_pos + 0.5

modules/custom_operations/examples/calculate_grid/export_model.py

@@ -0,0 +1,51 @@
+# Copyright (C) 2018-2022 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+import numpy as np
+import argparse
+import torch
+import torch.nn as nn
+from torch.autograd import Variable
+from .calculate_grid import CalculateGrid
+
+
+class MyModel(nn.Module):
+    def __init__(self):
+        super(MyModel, self).__init__()
+        self.calculate_grid = CalculateGrid()
+
+    def forward(self, x):
+        return self.calculate_grid.apply(x)
+
+
+def export(num_points, max_grid_extent):
+    # Generate a list of unique positions and add a mantissa
+    np.random.seed(32)
+    torch.manual_seed(11)
+
+    inp_pos = np.random.randint(0, max_grid_extent, [num_points, 3])
+    inp_pos = torch.tensor(inp_pos) + torch.rand(inp_pos.shape, dtype=torch.float32) # [0, 1)
+
+    model = MyModel()
+    with torch.no_grad():
+        torch.onnx.export(model, (inp_pos), 'model.onnx',
+                          input_names=['input'],
+                          output_names=['output'],
+                          operator_export_type=torch.onnx.OperatorExportTypes.ONNX_ATEN_FALLBACK)
+
+    ref = model(inp_pos).detach().numpy()
+
+    # Pad with sentinel values (-1, 0, 0) and zeros
+    ref = np.concatenate((ref, [[-1, 0, 0]]))
+    ref = np.pad(ref, ((0, inp_pos.shape[0] - ref.shape[0]), (0, 0)))
+
+    return [inp_pos.detach().numpy()], ref
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description='Generate ONNX model and test data')
+    parser.add_argument('--num_points', type=int, default=10)
+    parser.add_argument('--max_grid_extent', type=int, default=5)
+    args = parser.parse_args()
+
+    export(args.num_points, args.max_grid_extent)

modules/custom_operations/examples/complex_mul/complex_mul.py

@@ -0,0 +1,25 @@
+# Copyright (C) 2018-2022 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+import torch
+import torch.nn as nn
+
+class ComplexMul(torch.autograd.Function):
+    @staticmethod
+    def symbolic(g, input_tensor, other_tensor, is_conj = True):
+        return g.op("ComplexMultiplication", input_tensor, other_tensor, is_conj_i=int(is_conj))
+
+    @staticmethod
+    def forward(self, input_tensor, other_tensor):
+        complex_index = -1
+        real_part = input_tensor[..., 0] * other_tensor[..., 0] - input_tensor[..., 1] * other_tensor[..., 1]
+        imaginary_part = input_tensor[..., 0] * other_tensor[..., 1] + input_tensor[..., 1] * other_tensor[..., 0]
+
+        multiplication = torch.cat(
+            [
+                real_part.unsqueeze(dim=complex_index),
+                imaginary_part.unsqueeze(dim=complex_index),
+            ],
+            dim=complex_index,
+        )
+        return multiplication

modules/custom_operations/examples/complex_mul/export_model.py

@@ -0,0 +1,44 @@
+# Copyright (C) 2018-2022 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+import numpy as np
+import argparse
+import torch
+import torch.nn as nn
+from torch.autograd import Variable
+from .complex_mul import ComplexMul
+
+class MyModel(nn.Module):
+    def __init__(self):
+        super(MyModel, self).__init__()
+        self.complex_mul = ComplexMul()
+
+    def forward(self, x, y):
+        return self.complex_mul.apply(x, y)
+
+def export(inp_shape=[3, 2, 4, 8, 2], other_shape=[3, 2, 4, 8, 2]):
+    np.random.seed(324)
+    torch.manual_seed(32)
+
+    model = MyModel()
+    inp = Variable(torch.randn(inp_shape))
+    inp1 = Variable(torch.randn(other_shape))
+    model.eval()
+
+    with torch.no_grad():
+        torch.onnx.export(model, (inp, inp1), 'model.onnx',
+                        input_names=['input', 'input1'],
+                        output_names=['output'],
+                        operator_export_type=torch.onnx.OperatorExportTypes.ONNX_ATEN_FALLBACK)
+
+    ref = model(inp, inp1)
+    return [inp.detach().numpy(), inp1.detach().numpy()], ref.detach().numpy()
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description='Generate ONNX model and test data')
+    parser.add_argument('--inp_shape', type=int, nargs='+', default=[3, 2, 4, 8, 2])
+    parser.add_argument('--other_shape', type=int, nargs='+', default=[3, 2, 4, 8, 2])
+    args = parser.parse_args()
+
+    export(args.inp_shape, args.other_shape)

modules/custom_operations/examples/fft/export_model.py

@@ -0,0 +1,49 @@
+# Copyright (C) 2018-2022 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+import numpy as np
+import argparse
+import torch
+import torch.nn as nn
+from torch.autograd import Variable
+from .fft import FFT
+
+
+class MyModel(nn.Module):
+    def __init__(self, inverse, centered, dims):
+        super(MyModel, self).__init__()
+        self.inverse = inverse
+        self.centered = centered
+        self.dims = dims
+        self.fft = FFT()
+
+    def forward(self, x):
+        return self.fft.apply(x, self.inverse, self.centered, self.dims)
+
+
+def export(shape, inverse, centered, dims):
+    np.random.seed(324)
+    torch.manual_seed(32)
+
+    model = MyModel(inverse, centered, dims)
+    inp = Variable(torch.randn(shape))
+    model.eval()
+
+    with torch.no_grad():
+        torch.onnx.export(model, inp, 'model.onnx',
+                          input_names=['input'],
+                          output_names=['output'],
+                          operator_export_type=torch.onnx.OperatorExportTypes.ONNX_FALLTHROUGH)
+
+    ref = model(inp)
+    return [inp.detach().numpy()], ref.detach().numpy()
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description='Generate ONNX model and test data')
+    parser.add_argument('--shape', type=int, nargs='+', default=[5, 3, 6, 8, 2])
+    parser.add_argument('--inverse', action='store_true')
+    parser.add_argument('--centered', action='store_true')
+    parser.add_argument('--dims', type=int, nargs='+', default=[2, 3])
+    args = parser.parse_args()
+    export(args.shape, args.inverse, args.centered, args.dims)