OpenVINO refactoring

Author: dkurt

bonito/openvino/model.py

@@ -11,58 +11,24 @@
 except ImportError:
     pass
 
+
+def load_openvino_model(model, dirname):
+    package = model.config['model']['package']
+    if package == 'bonito.ctc':
+        return CTCModel(model, dirname)
+    elif package == 'bonito.crf':
+        return CRFModel(model, dirname)
+    else:
+        raise Exception('Unknown model configuration: ' + package)
+
+
 class OpenVINOModel:
 
-    def __init__(self, model, half, dirname):
+    def __init__(self, model):
         self.model = model
         self.alphabet = model.alphabet
         self.parameters = model.parameters
         self.stride = model.stride
-
-        package = model.config['model']['package']
-        if not package in ['bonito.ctc', 'bonito.crf']:
-            raise Exception('Unknown model configuration: ' + package)
-        self.is_ctc = package == 'bonito.ctc'
-        self.is_crf = package == 'bonito.crf'
-
-        if self.is_crf:
-            self.seqdist = model.seqdist
-            self.encoder = lambda data : self(data, encoder=True)
-
-        model_name = 'model' + ('_fp16' if half else '')
-        xml_path, bin_path = [os.path.join(dirname, model_name) + ext for ext in ['.xml', '.bin']]
-        self.ie = IECore()
-        if os.path.exists(xml_path) and os.path.exists(bin_path):
-            self.net = self.ie.read_network(xml_path, bin_path)
-        else:
-            # There is an issue with Swish at export step so we temporarly use default implementation
-            origin_swish_forward = Swish.forward
-            def swish_fake_forward(self, x):
-                return x * torch.sigmoid(x)
-            Swish.forward = swish_fake_forward
-
-            if self.is_crf:
-                inp = torch.randn([1, 1, 1000])
-                torch.onnx.export(model.encoder, inp, os.path.join(dirname, model_name) + '.onnx',
-                                  input_names=['input'], output_names=['output'],
-                                  opset_version=11)
-                raise Exception('OpenVINO 2021.2 is required to build CRF model in runtime. Use Model Optimizer instead.')
-
-            # Just a dummy input for export
-            inp = torch.randn([1, 1, 1, 1000])
-            buf = io.BytesIO()
-
-            # 1. Replace 1D layers to their 2D alternatives to improve efficiency
-            convert_to_2d(model)
-
-            # 2. Convert model to ONNX buffer
-            torch.onnx.export(model, inp, buf, input_names=['input'], output_names=['output'],
-                              opset_version=11)
-            Swish.forward = origin_swish_forward
-
-            # 3. Import network from memory buffer
-            self.net = self.ie.read_network(buf.getvalue(), b'', init_from_buffer=True)
-
         self.exec_net = None
 
 
@@ -78,10 +44,8 @@ def to(self, device):
         self.device = str(device).upper()
 
 
-    def __call__(self, data, encoder=False):
+    def process(self, data):
         data = data.float()
-        if self.is_ctc:
-            data = np.expand_dims(data, axis=2)  # 1D->2D
         batch_size = data.shape[0]
         inp_shape = list(data.shape)
         inp_shape[0] = 1  # We will run the batch asynchronously
@@ -135,15 +99,83 @@ def __call__(self, data, encoder=False):
             request = self.exec_net.requests[infer_request_id]
             output[:,out_id:out_id+1] = request.output_blobs['output'].buffer
 
-        output = torch.tensor(output)
-        if encoder:
-            return output
-        return self.model.global_norm(output) if self.is_crf else output
+        return torch.tensor(output)
+
+
+class CTCModel(OpenVINOModel):
+
+    def __init__(self, model, dirname):
+        super().__init__(model)
+
+        model_name = 'model'
+        xml_path, bin_path = [os.path.join(dirname, model_name) + ext for ext in ['.xml', '.bin']]
+        self.ie = IECore()
+        if os.path.exists(xml_path) and os.path.exists(bin_path):
+            self.net = self.ie.read_network(xml_path, bin_path)
+        else:
+            # There is an issue with Swish at export step so we temporarly use default implementation
+            origin_swish_forward = Swish.forward
+            def swish_fake_forward(self, x):
+                return x * torch.sigmoid(x)
+            Swish.forward = swish_fake_forward
+
+            # Just a dummy input for export
+            inp = torch.randn([1, 1, 1, 1000])
+            buf = io.BytesIO()
+
+            # 1. Replace 1D layers to their 2D alternatives to improve efficiency
+            convert_to_2d(model)
+
+            # 2. Convert model to ONNX buffer
+            torch.onnx.export(model, inp, buf, input_names=['input'], output_names=['output'],
+                              opset_version=11)
+            Swish.forward = origin_swish_forward
+
+            # 3. Import network from memory buffer
+            self.net = self.ie.read_network(buf.getvalue(), b'', init_from_buffer=True)
+
+
+    def __call__(self, data):
+        data = data.unsqueeze(2)  # 1D->2D
+        return self.process(data)
 
 
     def decode(self, x, beamsize=5, threshold=1e-3, qscores=False, return_path=False):
-        if self.is_crf:
-            return self.model.decode(x)
+        return self.model.decode(x, beamsize=beamsize, threshold=threshold,
+                                 qscores=qscores, return_path=return_path)
+
+
+class CRFModel(OpenVINOModel):
+
+    def __init__(self, model, dirname):
+        super().__init__(model)
+        self.seqdist = model.seqdist
+        self.encoder = lambda data : self(data, encoder=True)
+
+        # TODO: move to OpenVINOModel constructor when 2021.2 is out
+        model_name = 'model'
+        xml_path, bin_path = [os.path.join(dirname, model_name) + ext for ext in ['.xml', '.bin']]
+        self.ie = IECore()
+        if os.path.exists(xml_path) and os.path.exists(bin_path):
+            self.net = self.ie.read_network(xml_path, bin_path)
         else:
-            return self.model.decode(x, beamsize=beamsize, threshold=threshold,
-                                     qscores=qscores, return_path=return_path)
+            inp = torch.randn([1, 1, 1000])
+            torch.onnx.export(model.encoder, inp, os.path.join(dirname, model_name) + '.onnx',
+                              input_names=['input'], output_names=['output'],
+                              opset_version=11)
+            raise Exception('OpenVINO 2021.2 is required to build CRF model in runtime. Use Model Optimizer instead.')
+
+
+    def __call__(self, data, encoder=False):
+        if encoder:
+            return self.process(data)
+        else:
+            return self.model.global_norm(self.process(data))
+
+
+    def decode(self, x):
+        return self.model.decode(x)
+
+
+    def decode_batch(self, x):
+        return self.model.decode_batch(x)

bonito/util.py

@@ -17,7 +17,7 @@
 import parasail
 import numpy as np
 from torch.cuda import get_device_capability
-from bonito.openvino.model import OpenVINOModel
+from bonito.openvino.model import load_openvino_model
 
 try:
     from claragenomics.bindings import cuda
@@ -295,7 +295,7 @@ def load_model(dirname, device, weights=None, half=None, chunksize=0, use_openvi
     model.load_state_dict(new_state_dict)
 
     if use_openvino:
-        model = OpenVINOModel(model, half, dirname)
+        model = load_openvino_model(model, dirname)
 
     if half is None and device != torch.device('cpu'):
         half = half_supported()