Use OpenVINO 2021.2

README.md

@@ -33,7 +33,6 @@ $ source venv3/bin/activate
 To optimize inference on CPU with Intel OpenVINO:
 
 ```bash
-(venv3) $ export LD_LIBRARY_PATH=$(pwd)/venv3/lib:$LD_LIBRARY_PATH
 (venv3) $ bonito evaluate dna_r9.4.1 --use_openvino --device=cpu
 ```
 

bonito/openvino/model.py

@@ -24,18 +24,15 @@ def load_openvino_model(model, dirname):
 
 class OpenVINOModel:
 
-    def __init__(self, model):
+    def __init__(self, model, dirname):
         self.model = model
         self.alphabet = model.alphabet
         self.parameters = model.parameters
         self.stride = model.stride
+        self.net = None
         self.exec_net = None
-
-        # 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
+        self.dirname = dirname
+        self.ie = IECore()
 
 
     def eval(self):
@@ -49,21 +46,44 @@ def half(self):
     def to(self, device):
         self.device = str(device).upper()
 
+    """
+    Call this method once to initialize executable network
+    """
+    def init_model(self, model, inp_shape):
+        # First, we try to check if there is IR on disk. If not - load model in runtime
+        xml_path, bin_path = [os.path.join(self.dirname, 'model') + ext for ext in ['.xml', '.bin']]
+        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
+
+            # Convert model to ONNX buffer
+            buf = io.BytesIO()
+            inp = torch.randn(inp_shape)
+            torch.onnx.export(model, inp, buf, input_names=['input'], output_names=['output'],
+                              opset_version=11)
+            Swish.forward = origin_swish_forward
+
+            # Import network from memory buffer
+            self.net = self.ie.read_network(buf.getvalue(), b'', init_from_buffer=True)
+
+        # Load model to device
+        config = {}
+        if self.device == 'CPU':
+            config={'CPU_THROUGHPUT_STREAMS': 'CPU_THROUGHPUT_AUTO'}
+        self.exec_net = self.ie.load_network(self.net, self.device,
+                                             config=config, num_requests=0)
+
 
     def process(self, data):
         data = data.float()
         batch_size = data.shape[0]
         inp_shape = list(data.shape)
         inp_shape[0] = 1  # We will run the batch asynchronously
-        if self.net.input_info['input'].tensor_desc.dims != inp_shape:
-            self.net.reshape({'input': inp_shape})
-            self.exec_net = None
-        if not self.exec_net:
-            config = {}
-            if self.device == 'CPU':
-                config={'CPU_THROUGHPUT_STREAMS': 'CPU_THROUGHPUT_AUTO'}
-            self.exec_net = self.ie.load_network(self.net, self.device,
-                                                 config=config, num_requests=0)
 
         # List that maps infer requests to index of processed chunk from batch.
         # -1 means that request has not been started yet.
@@ -111,32 +131,15 @@ def process(self, data):
 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:
-            # 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)
+        super().__init__(model, dirname)
 
 
     def __call__(self, data):
         data = data.unsqueeze(2)  # 1D->2D
+        if self.exec_net is None:
+            convert_to_2d(self.model)
+            self.init_model(self.model, [1, 1, 1, data.shape[-1]])
+
         return self.process(data)
 
 
@@ -148,25 +151,15 @@ def decode(self, x, beamsize=5, threshold=1e-3, qscores=False, return_path=False
 class CRFModel(OpenVINOModel):
 
     def __init__(self, model, dirname):
-        super().__init__(model)
+        super().__init__(model, dirname)
         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:
-            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 self.exec_net is None:
+            self.init_model(self.model.encoder, [1, 1, data.shape[-1]])
+
         if encoder:
             return self.process(data)
         else:

requirements.txt

@@ -1,7 +1,7 @@
 mappy==2.17
 toml==0.10.0
 tqdm==4.31.1
-numpy<=1.16.3
+numpy<=1.18.5
 torch>=1.1.0,<=1.5
 optuna==1.1.0
 parasail==1.2
@@ -13,5 +13,4 @@ ont-fast5-api==3.1.6
 fast-ctc-decode==0.2.5
 #bonito-cuda-runtime==0.0.2a2
 #pyclaragenomics-cuda-10-0==0.4.2
-tbb==2020.3.254
-openvino-python==2021.1
+experimental-openvino-python==2021.2