Build OpenVINO network in runtime

Author: dkurt

README.md

@@ -79,15 +79,6 @@ $ source venv3/bin/activate
 
 To optimize inference on CPU with Intel OpenVINO:
 
-1. Convert model to OpenVINO IR (optional if you already have `.xml` and `.bin` files)
-```bash
-(venv3) $ git clone -b releases/2021/1 --depth 1 https://github.com/openvinotoolkit/openvino
-(venv3) $ export PYTHONPATH=openvino/model-optimizer/:$PYTHONPATH
-(venv3) $ pip install -r openvino/model-optimizer/requirements_onnx.txt
-(venv3) $ python bonito/openvino/convert.py dna_r9.4.1
-```
-
-2. Run evaluation
 ```bash
 (venv3) $ export LD_LIBRARY_PATH=$(pwd)/venv3/lib:$LD_LIBRARY_PATH
 (venv3) $ bonito evaluate dna_r9.4.1 --use_openvino --device=cpu

bonito/model.py

@@ -50,6 +50,11 @@ def forward(self, x):
         return SwishAutoFn.apply(x)
 
 
+class Sum(Module):
+    def forward(self, x, y):
+        return x + y
+
+
 activations = {
     "relu": ReLU,
     "swish": Swish,
@@ -166,6 +171,7 @@ def __init__(self, in_channels, out_channels, activation, repeat=5, kernel_size=
 
         self.use_res = residual
         self.conv = ModuleList()
+        self.sum = Sum()
 
         _in_channels = in_channels
         padding = self.get_padding(kernel_size[0], stride[0], dilation[0])
@@ -223,7 +229,7 @@ def forward(self, x):
         for layer in self.conv:
             _x = layer(_x)
         if self.use_res:
-            _x += self.residual(x)
+            _x = self.sum(_x, self.residual(x))
         return self.activation(_x)
 
 

bonito/openvino/convert.py

@@ -0,0 +1,100 @@
+# This script provides a method which builds OpenVINO network in runtime
+import numpy as np
+from openvino.inference_engine import IECore, IENetwork
+
+import ngraph.opset4 as ng
+from ngraph.impl.op import Parameter
+from ngraph.impl import Function, Shape, Type
+
+import torch
+from torch.autograd import Variable
+
+
+nodes = {}
+out = None
+
+def forward_hook(self, inputs, output):
+    global out
+    layer_type = self.__class__.__name__
+
+    params = [value.numpy() for value in self.state_dict().values()]
+
+    inp = nodes[inputs[0].data_ptr()]
+    if layer_type == 'Conv1d':
+        weights = np.expand_dims(params[0], axis=2)
+        if self.groups == 1:
+            out = ng.convolution(inp, weights,
+                                 [1, self.stride[0]],
+                                 [0, self.padding[0]],
+                                 [0, self.padding[0]],
+                                 [1, self.dilation[0]])
+
+        else:
+            weights = weights.reshape(self.groups, weights.shape[0] // self.groups, weights.shape[1], weights.shape[2], weights.shape[3])
+            out = ng.group_convolution(inp, weights,
+                                       [1, self.stride[0]],
+                                       [0, self.padding[0]],
+                                       [0, self.padding[0]],
+                                       [1, self.dilation[0]])
+        if len(params) > 1:
+            assert(len(params) == 2)
+            bias = params[1].reshape(1, params[1].shape[0], 1, 1)
+            out = ng.add(out, bias)
+
+    elif layer_type == 'BatchNorm1d':
+        out = ng.batch_norm_inference(inp, params[0], params[1], params[2], params[3], self.eps)
+    elif layer_type == 'Swish':
+        out = ng.swish(inp)
+    elif layer_type == 'Sum':
+        out = ng.add(inp, nodes[inputs[1].data_ptr()])
+    elif layer_type == 'Dropout':
+        return
+    else:
+        raise Exception('Unknown layer type: ' + layer_type)
+
+    nodes[output.data_ptr()] = out
+
+
+def sanity_check(net, inp, ref):
+    ie = IECore()
+    exec_net = ie.load_network(net, 'CPU')
+    ie_out = exec_net.infer({'input': inp.numpy()})
+    ie_out = next(iter(ie_out.values()))
+
+    ref = ref.numpy().reshape(ie_out.shape)
+    diff = np.max(np.abs(ie_out - ref))
+    print('PyTorch / OpenVINO diff:', diff)
+    print('Reference values range: [{}, {}]'.format(np.min(ref), np.max(ref)))
+    if diff > 1.1e-4:
+        raise Exception('Sanity check failed with diff', diff)
+
+
+def torch2openvino(model):
+    with torch.no_grad():
+        model.eval()
+        for module in model.modules():
+            if len([m for m in module.modules()]) != 1:
+                continue
+            module.register_forward_hook(forward_hook)
+
+        # Just a dummy input to make forward pass
+        inp = Variable(torch.randn([1, 1, 1000]))
+
+        param = Parameter(Type.f32, Shape([1, 1, 1, 1000]))
+        nodes[inp.data_ptr()] = param
+        ref = model(inp)
+
+    out_node = ng.transpose(out, [0, 3, 2, 1])
+    out_node = ng.log(ng.softmax(out_node, axis=3))
+
+    param.set_friendly_name('input')
+    out_node.set_friendly_name('output')
+    func = Function([out_node], [param], '')
+
+    caps = Function.to_capsule(func)
+    net = IENetwork(caps)
+
+    # Uncomment to perform conversion check
+    # sanity_check(net, inp, ref)
+
+    return net

bonito/openvino/loader.py

@@ -4,6 +4,7 @@
 
 try:
     from openvino.inference_engine import IECore, StatusCode
+    from .loader import torch2openvino
 except ImportError:
     pass
 
@@ -16,7 +17,10 @@ def __init__(self, model, half, dirname):
         model_name = model.config['model'] + ('_fp16' if half else '')
         xml_path, bin_path = [os.path.join(dirname, model_name) + ext for ext in ['.xml', '.bin']]
         self.ie = IECore()
-        self.net = self.ie.read_network(xml_path, bin_path)
+        if os.path.exists(xml_path) and os.path.exists(bin_path):
+            self.net = self.ie.read_network(xml_path, bin_path)
+        else:
+            self.net = torch2openvino(model)
         self.exec_net = None