Use test pip image for OpenVINO

Author: dkurt

README.md

@@ -78,10 +78,18 @@ $ source venv3/bin/activate
 ```
 
 To optimize inference on CPU with Intel OpenVINO:
-* Download and install OpenVINO from https://software.seek.intel.com/openvino-toolkit
+
+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) $ pip install -r /opt/intel/openvino/deployment_tools/model_optimizer/requirements_onnx.txt
-(venv3) $ source /opt/intel/openvino/bin/setupvars.sh
+(venv3) $ export LD_LIBRARY_PATH=$(pwd)/venv3/lib:$LD_LIBRARY_PATH
 (venv3) $ bonito evaluate dna_r9.4.1 --use_openvino --device=cpu
 ```
 

bonito/openvino/convert.py

@@ -0,0 +1,43 @@
+import os
+import torch
+import argparse
+from bonito.util import load_model
+
+parser = argparse.ArgumentParser(description='Use this script to prepare OpenVINO IR for trained model')
+parser.add_argument("model_directory")
+parser.add_argument("--half", action="store_true", default=False)
+parser.add_argument("--weights", default="0", type=str)
+args = parser.parse_args()
+
+__dir__ = os.path.dirname(os.path.dirname(__file__))
+__models__ = os.path.join(__dir__, "models")
+dirname = args.model_directory
+if not os.path.isdir(dirname) and os.path.isdir(os.path.join(__models__, dirname)):
+    dirname = os.path.join(__models__, dirname)
+
+model = load_model(dirname, 'cpu', weights=int(args.weights), half=args.half)
+
+
+# Convert to ONNX
+onnx_path = os.path.join(dirname, model.config['model']) + '.onnx'
+inp = torch.randn(1, 1, 1000)  # Just dummy input shape. We will reshape model later
+model.eval()
+with torch.no_grad():
+    torch.onnx.export(model, inp, onnx_path,
+                    input_names=['input'],
+                    output_names=['output'],
+                    operator_export_type=torch.onnx.OperatorExportTypes.ONNX_ATEN_FALLBACK)
+
+
+# Convert to IR
+import mo_onnx
+import subprocess
+model_name = model.config['model'] + ('_fp16' if args.half else '')
+subprocess.call([mo_onnx.__file__,
+                 '--input_model', onnx_path,
+                 '--extension', os.path.join(os.path.dirname(__file__), 'mo_extension'),
+                 '--keep_shape_ops',
+                 '--model_name', model_name,
+                 '--data_type', 'FP16' if args.half else 'FP32',
+                 '--input_shape=[1,1,1,1000]',
+                 '--output_dir', dirname])

bonito/openvino/model.py

@@ -13,33 +13,8 @@ def __init__(self, model, half, dirname):
         self.model = model
         self.alphabet = model.alphabet
 
-        onnx_path = os.path.join(dirname, model.config['model']) + '.onnx'
         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']]
-        if not os.path.exists(xml_path) or not os.path.exists(bin_path):
-
-            # Convert to ONNX
-            if not os.path.exists(onnx_path):
-                inp = torch.randn(1, 1, 1000)  # Just dummy input shape. We will reshape model later
-                model.eval()
-                with torch.no_grad():
-                    torch.onnx.export(model, inp, onnx_path,
-                                    input_names=['input'],
-                                    output_names=['output'],
-                                    operator_export_type=torch.onnx.OperatorExportTypes.ONNX_ATEN_FALLBACK)
-
-            # Convert to IR
-            import mo_onnx
-            import subprocess
-            subprocess.call([mo_onnx.__file__,
-                             '--input_model', onnx_path,
-                             '--extension', os.path.join(os.path.dirname(__file__), 'mo_extension'),
-                             '--keep_shape_ops',
-                             '--model_name', model_name,
-                             '--data_type', 'FP16' if half else 'FP32',
-                             '--input_shape=[1,1,1,1000]',
-                             '--output_dir', dirname])
-
         self.ie = IECore()
         self.net = self.ie.read_network(xml_path, bin_path)
         self.exec_net = None

requirements.txt

@@ -9,3 +9,5 @@ ont-fast5-api==3.0.1
 fast-ctc-decode==0.2.5
 bonito-cuda-runtime==0.0.2a2
 #pyclaragenomics-cuda-10-0==0.4.2
+--extra-index-url https://test.pypi.org/simple/
+openvino

setup.py

@@ -13,7 +13,10 @@
     raise RuntimeError('Unable to find version string in "{}/__init__.py".'.format(__pkg_name__))
 
 with open('requirements.txt') as f:
-    requirements = f.read().splitlines()
+    requirements = []
+    for req in f.read().splitlines():
+        if not req.startswith('--extra-index-url'):
+            requirements.append(req)
 
 with open('README.md', encoding='utf-8') as f:
     long_description = f.read()