Intel OpenVINO backend

.github/workflows/build.yml

@@ -40,6 +40,10 @@ jobs:
         key: ${{ matrix.os }}-pip-${{ hashFiles('env.*.yml') }}
         restore-keys: |
           ${{ runner.os }}-pip-
+    - name: Set up Python ${{ matrix.python-version }}
+      uses: actions/setup-python@v2
+      with:
+        python-version: ${{ matrix.python-version }}
     - name: Install miniconda
       uses: conda-incubator/setup-miniconda@v2
       with:
@@ -65,6 +69,7 @@ jobs:
         python -c 'import rdkit; print(rdkit.__version__)'
         python -c 'import torch; print(torch.__version__)'
         python -c 'import tensorflow; print(tensorflow.__version__)'
+        python -c 'import openvino.inference_engine as ie; print(ie.__version__)'
 
   powershell-build:
     runs-on: windows-latest
@@ -104,3 +109,4 @@ jobs:
         python -c 'import rdkit; print(rdkit.__version__)'
         python -c 'import torch; print(torch.__version__)'
         python -c 'import tensorflow; print(tensorflow.__version__)'
+        python -c 'import openvino.inference_engine as ie; print(ie.__version__)'

deepchem/models/keras_model.py

@@ -15,6 +15,7 @@
 from deepchem.models.optimizers import Adam, Optimizer, LearningRateSchedule
 from deepchem.trans import Transformer, undo_transforms
 from deepchem.utils.evaluate import GeneratorEvaluator
+from deepchem.utils.openvino_model import OpenVINOModel
 
 from typing import Any, Callable, Dict, Iterable, List, Optional, Sequence, Tuple, Union
 from deepchem.utils.typing import ArrayLike, LossFn, OneOrMany
@@ -226,6 +227,11 @@ def __init__(self,
     self._output_functions: Dict[Any, Any] = {}
     self._gradient_fn_for_vars: Dict[Any, Any] = {}
 
+    self._use_openvino = kwargs.get('use_openvino', False)
+    if self._use_openvino:
+      self._openvino_model = OpenVINOModel(
+          self.model_dir, self.batch_size, keras_model=self)
+
   def _ensure_built(self) -> None:
     """The first time this is called, create internal data structures."""
     if self._built:
@@ -595,6 +601,10 @@ def _predict(
       )
     if tf.is_tensor(outputs):
       outputs = [outputs]
+
+    if self._use_openvino:
+      openvino_predictions, generator = self._openvino_model(generator)
+
     for batch in generator:
       inputs, labels, weights = batch
       self._create_inputs(inputs)
@@ -611,10 +621,14 @@ def _predict(
               self.model.inputs, outputs)
         output_values = self._output_functions[key](inputs)
       else:
-        output_values = self._compute_model(inputs)
-        if tf.is_tensor(output_values):
+        if self._use_openvino:
+          output_values = next(openvino_predictions)
           output_values = [output_values]
-        output_values = [t.numpy() for t in output_values]
+        else:
+          output_values = self._compute_model(inputs)
+          if tf.is_tensor(output_values):
+            output_values = [output_values]
+          output_values = [t.numpy() for t in output_values]
 
       # Apply tranformers and record results.
       if uncertainty:

deepchem/models/tests/test_openvino.py

@@ -0,0 +1,105 @@
+import pytest
+import numpy as np
+import tensorflow as tf
+import torch
+import os
+
+import deepchem as dc
+from deepchem.molnet import load_tox21, load_perovskite
+from deepchem.data import NumpyDataset
+from deepchem.models.torch_models import TorchModel
+
+
+def test_cgcnn_regression():
+  np.random.seed(123)
+  torch.manual_seed(124)
+
+  # load datasets
+  current_dir = os.path.dirname(os.path.abspath(__file__))
+  config = {
+      "reload": False,
+      "featurizer": dc.feat.CGCNNFeaturizer(),
+      # disable transformer
+      "transformers": [],
+      "data_dir": current_dir
+  }
+
+  tasks, datasets, transformers = load_perovskite(**config)
+  train, valid, _ = datasets
+  n_tasks = len(tasks)
+
+  def init(**kwargs):
+    torch.manual_seed(124)
+    return dc.models.CGCNNModel(
+        n_tasks=n_tasks,
+        mode='regression',
+        batch_size=1,
+        learning_rate=0.001,
+        **kwargs)
+
+  ref_model = init()
+  model = init(use_openvino=True)
+
+  # check overfit
+  metric = dc.metrics.Metric(dc.metrics.mae_score, n_tasks=n_tasks)
+  ref_scores = ref_model.evaluate(train, [metric], transformers)
+  scores = model.evaluate(train, [metric], transformers)
+
+  if os.path.exists(os.path.join(current_dir, 'perovskite.json')):
+    os.remove(os.path.join(current_dir, 'perovskite.json'))
+
+  assert abs(scores['mae_score'] - ref_scores['mae_score']) < 0.2
+
+
+def test_tox21_tf_progressive():
+  np.random.seed(123)
+  tf.random.set_seed(124)
+
+  n_features = 1024
+  tox21_tasks, tox21_datasets, transformers = load_tox21()
+  _, valid_dataset, _ = tox21_datasets
+
+  metric = dc.metrics.Metric(dc.metrics.roc_auc_score, np.mean)
+
+  def init(**kwargs):
+    tf.random.set_seed(572)
+    return dc.models.ProgressiveMultitaskClassifier(
+        len(tox21_tasks),
+        n_features,
+        layer_sizes=[1000],
+        dropouts=[.25],
+        learning_rate=0.001,
+        batch_size=50,
+        **kwargs)
+
+  ref_model = init()
+  model = init(use_openvino=True)
+
+  ref_scores = ref_model.evaluate(valid_dataset, [metric], transformers)
+  scores = model.evaluate(valid_dataset, [metric], transformers)
+
+  assert scores['mean-roc_auc_score'] == pytest.approx(
+      ref_scores['mean-roc_auc_score'], 1e-5)
+  assert model._openvino_model.is_available()
+
+
+def test_torch_model():
+  np.random.seed(123)
+
+  def init(**kwargs):
+    torch.manual_seed(1412)
+    pytorch_model = torch.nn.Sequential(
+        torch.nn.Linear(1024, 1000), torch.nn.ReLU(), torch.nn.Dropout(0.5),
+        torch.nn.Linear(1000, 1))
+    return TorchModel(pytorch_model, dc.models.losses.L2Loss(), **kwargs)
+
+  ref_model = init(batch_size=8)
+  model = init(batch_size=8, use_openvino=True)
+
+  dataset = NumpyDataset(np.random.standard_normal((100, 1024)))
+  ref = ref_model.predict(dataset)
+  out = model.predict(dataset)
+  diff = np.abs(np.max(ref - out))
+
+  assert diff < 1e-6
+  assert model._openvino_model.is_available()

deepchem/models/torch_models/torch_model.py

@@ -16,6 +16,7 @@
 from deepchem.models.optimizers import Adam, Optimizer, LearningRateSchedule
 from deepchem.trans import Transformer, undo_transforms
 from deepchem.utils.evaluate import GeneratorEvaluator
+from deepchem.utils.openvino_model import OpenVINOModel
 
 from typing import Any, Callable, Dict, Iterable, List, Optional, Sequence, Tuple, Union
 from deepchem.utils.typing import ArrayLike, LossFn, OneOrMany
@@ -216,6 +217,11 @@ def __init__(self,
     self._output_functions: Dict[Any, Any] = {}
     self._optimizer_for_vars: Dict[Any, Any] = {}
 
+    self._use_openvino = kwargs.get('use_openvino', False)
+    if self._use_openvino:
+      self._openvino_model = OpenVINOModel(
+          self.model_dir, self.batch_size, torch_model=self)
+
   def _ensure_built(self) -> None:
     """The first time this is called, create internal data structures."""
     if self._built:
@@ -520,14 +526,21 @@ def _predict(
         )
     self._ensure_built()
     self.model.eval()
+
+    if self._use_openvino:
+      openvino_predictions, generator = self._openvino_model(generator)
+
     for batch in generator:
       inputs, labels, weights = batch
       inputs, _, _ = self._prepare_batch((inputs, None, None))
 
       # Invoke the model.
       if isinstance(inputs, list) and len(inputs) == 1:
         inputs = inputs[0]
-      output_values = self.model(inputs)
+      if self._use_openvino:
+        output_values = [torch.Tensor(next(openvino_predictions))]
+      else:
+        output_values = self.model(inputs)
       if isinstance(output_values, torch.Tensor):
         output_values = [output_values]
       output_values = [t.detach().cpu().numpy() for t in output_values]