OpenVINO optimizations for Knowledge graphs

This tutorial is also available as a Jupyter notebook that can be cloned directly from GitHub. See the installation guide for instructions to run this tutorial locally on Windows, Linux or macOS.

Github

The goal of this notebook is to showcase performance optimizations for the ConvE knowledge graph embeddings model using the Intel® Distribution of OpenVINO™ Toolkit. The optimizations process contains the following steps: 1. Export the trained model to a format suitable for OpenVINO optimizations and inference 2. Report the inference performance speedup obtained with the optimized OpenVINO model

The ConvE model we use is an implementation of the paper Convolutional 2D Knowledge Graph Embeddings (https://arxiv.org/abs/1707.01476). The sample dataset was downloaded from: https://github.com/TimDettmers/ConvE/tree/master/countries/countries_S1

Windows specific settings

# On Windows, add the directory that contains cl.exe to the PATH
# to enable PyTorch to find the required C++ tools.
# This code assumes that Visual Studio 2019 is installed in the default directory.
# If you have a different C++ compiler, please add the correct path
# to os.environ["PATH"] directly.
# Note that the C++ Redistributable is not enough to run this notebook.

# Adding the path to os.environ["LIB"] is not always required
# - it depends on the system's configuration

import sys

if sys.platform == "win32":
    import distutils.command.build_ext
    import os
    from pathlib import Path

    VS_INSTALL_DIR = r"C:/Program Files (x86)/Microsoft Visual Studio"
    cl_paths = sorted(list(Path(VS_INSTALL_DIR).glob("**/Hostx86/x64/cl.exe")))
    if len(cl_paths) == 0:
        raise ValueError(
            "Cannot find Visual Studio. This notebook requires a C++ compiler. If you installed "
            "a C++ compiler, please add the directory that contains"
            "cl.exe to `os.environ['PATH']`."
        )
    else:
        # If multiple versions of MSVC are installed, get the most recent version
        cl_path = cl_paths[-1]
        vs_dir = str(cl_path.parent)
        os.environ["PATH"] += f"{os.pathsep}{vs_dir}"
        # Code for finding the library dirs from
        # https://stackoverflow.com/questions/47423246/get-pythons-lib-path
        d = distutils.core.Distribution()
        b = distutils.command.build_ext.build_ext(d)
        b.finalize_options()
        os.environ["LIB"] = os.pathsep.join(b.library_dirs)
        print(f"Added {vs_dir} to PATH")

Import the packages needed for successful execution

import numpy as np
import time
import json

import torch
from torch.nn import functional as F, Parameter
from torch.nn.init import xavier_normal_

from pathlib import Path

from sklearn.metrics import accuracy_score

from openvino.runtime import Core

Settings: Including path to the serialized model files and input data files

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using {device} device")

# Path to the trained model
modelpath = Path('models/conve.pt')

# Path to the file containing the entities and entity IDs
entdatapath = Path('data/countries_S1/kg_training_entids.txt')

# Path to the file containing the relations and relation IDs
reldatapath = Path('data/countries_S1/kg_training_relids.txt')

# Path to the test data file
testdatapath = Path('data/countries_S1/e1rel_to_e2_ranking_test.json')

# Entity and relation embedding dimensions
EMB_DIM = 300

# Top K vals to consider from the predictions
TOP_K = 2

# Required for OpenVINO conversion
output_dir = Path("models")
base_model_name = "conve"

output_dir.mkdir(exist_ok=True)

# Paths where PyTorch, ONNX and OpenVINO IR models will be stored
fp32_onnx_path = Path(output_dir / (base_model_name + "_fp32")).with_suffix(".onnx")
Using cpu device

Defining the ConvE model class

# Model implementation reference: https://github.com/TimDettmers/ConvE
class ConvE(torch.nn.Module):
    def __init__(self, num_entities, num_relations, emb_dim):
        super(ConvE, self).__init__()
        # Embedding tables for entity and relations with num_uniq_ent in y-dim, emb_dim in x-dim
        self.emb_e = torch.nn.Embedding(num_entities, emb_dim, padding_idx=0)
        self.ent_weights_matrix = torch.ones([num_entities, emb_dim], dtype=torch.float64)
        self.emb_rel = torch.nn.Embedding(num_relations, emb_dim, padding_idx=0)
        self.ne = num_entities
        self.nr = num_relations
        self.inp_drop = torch.nn.Dropout(0.2)
        self.hidden_drop = torch.nn.Dropout(0.3)
        self.feature_map_drop = torch.nn.Dropout2d(0.2)
        self.loss = torch.nn.BCELoss()
        self.conv1 = torch.nn.Conv2d(1, 32, (3, 3), 1, 0, bias=True)
        self.bn0 = torch.nn.BatchNorm2d(1)
        self.bn1 = torch.nn.BatchNorm2d(32)
        self.ln0 = torch.nn.LayerNorm(emb_dim)
        self.register_parameter('b', Parameter(torch.zeros(num_entities)))
        self.fc = torch.nn.Linear(16128, emb_dim)

    def init(self):
        """ Initializes the model """
        # Xavier initialization
        xavier_normal_(self.emb_e.weight.data)
        xavier_normal_(self.emb_rel.weight.data)

    def forward(self, e1, rel):
        """ Forward pass on the model.
        :param e1: source entity
        :param rel: relation between the source and target entities
        Returns the model predictions for the target entities
        """
        e1_embedded = self.emb_e(e1).view(-1, 1, 10, 30)
        rel_embedded = self.emb_rel(rel).view(-1, 1, 10, 30)
        stacked_inputs = torch.cat([e1_embedded, rel_embedded], 2)
        stacked_inputs = self.bn0(stacked_inputs)
        x = self.inp_drop(stacked_inputs)
        x = self.conv1(x)
        x = self.bn1(x)
        x = F.relu(x)
        x = self.feature_map_drop(x)
        x = x.view(1, -1)
        x = self.fc(x)
        x = self.hidden_drop(x)
        x = self.ln0(x)
        x = F.relu(x)
        x = torch.mm(x, self.emb_e.weight.transpose(1, 0))
        x = self.hidden_drop(x)
        x += self.b.expand_as(x)
        pred = torch.nn.functional.softmax(x, dim=1)
        return pred

Defining the dataloader

class DataLoader():
    def __init__(self):
        super(DataLoader, self).__init__()

        self.ent_path = entdatapath
        self.rel_path = reldatapath
        self.test_file = testdatapath
        self.entity_ids, self.ids2entities = self.load_data(data_path=self.ent_path)
        self.rel_ids, self.ids2rel = self.load_data(data_path=self.rel_path)
        self.test_triples_list = self.convert_triples(data_path=self.test_file)

    def load_data(self, data_path):
        """ Creates a dictionary of data items with corresponding ids """
        item_dict, ids_dict = {}, {}
        fp = open(data_path, "r")
        lines = fp.readlines()
        for line in lines:
            name, id = line.strip().split('\t')
            item_dict[name] = int(id)
            ids_dict[int(id)] = name
        fp.close()
        return item_dict, ids_dict

    def convert_triples(self, data_path):
        """ Creates a triple of source entity, relation and target entities"""
        triples_list = []
        dp = open(data_path, "r")
        lines = dp.readlines()
        for line in lines:
            item_dict = json.loads(line.strip())
            h = item_dict['e1']
            r = item_dict['rel']
            t = item_dict['e2_multi1'].split('\t')
            hrt_list = []
            hrt_list.append(self.entity_ids[h])
            hrt_list.append(self.rel_ids[r])
            t_ents = []
            for t_idx in t:
                t_ents.append(self.entity_ids[t_idx])
            hrt_list.append(t_ents)
            triples_list.append(hrt_list)
        dp.close()
        return triples_list

Evaluate the trained ConvE model

We will first evaluate the model performance using PyTorch. The goal is to make sure there are no accuracy differences between the original model inference and the model converted to OpenVINO intermediate representation inference results. Here, we use a simple accuracy metric to evaluate the model performance on a test dataset. However, it is typical to use metrics such as Mean Reciprocal Rank, Hits@10 etc.

data = DataLoader()
num_entities = len(data.entity_ids)
num_relations = len(data.rel_ids)

model = ConvE(num_entities=num_entities, num_relations=num_relations, emb_dim=EMB_DIM)
model.load_state_dict(torch.load(modelpath))
model.eval()

pt_inf_times = []

triples_list = data.test_triples_list
num_test_samples = len(triples_list)
pt_acc = 0.0
for i in range(num_test_samples):
    test_sample = triples_list[i]
    h, r, t = test_sample
    start_time = time.time()
    logits = model.forward(e1=torch.tensor(h), rel=torch.tensor(r))
    end_time = time.time()
    pt_inf_times.append(end_time - start_time)
    score, pred = torch.topk(logits, TOP_K, 1)

    gt = np.array(sorted(t))
    pred = np.array(sorted(pred[0].cpu().detach()))
    pt_acc += accuracy_score(gt, pred)

avg_pt_time = np.mean(pt_inf_times) * 1000
print(f'Average time taken for inference: {avg_pt_time} ms')
print(f'Mean accuracy of the model on the test dataset: {pt_acc/num_test_samples}')
Average time taken for inference: 0.6341437498728434 ms
Mean accuracy of the model on the test dataset: 0.875

Prediction on the Knowledge graph.

As a sample evaluation, we perform the entity prediction task on the knowledge graph. We pass the source entity ‘san_marino’ and relation ‘locatedIn’ to the knowledge graph and obtain the target entity predictions. We expect to see as predictions, target entities that form a factual triple with the entity and relation passed as inputs to the knowledge graph.

entitynames_dict = data.ids2entities

ent = 'san_marino'
rel = 'locatedin'

h_idx = data.entity_ids[ent]
r_idx = data.rel_ids[rel]

logits = model.forward(torch.tensor(h_idx), torch.tensor(r_idx))
score, pred = torch.topk(logits, TOP_K, 1)

for j, id in enumerate(pred[0].cpu().detach().numpy()):
    pred_entity = entitynames_dict[id]
    print(f'Source Entity: {ent}, Relation: {rel}, Target entity prediction: {pred_entity}')
Source Entity: san_marino, Relation: locatedin, Target entity prediction: southern_europe
Source Entity: san_marino, Relation: locatedin, Target entity prediction: europe

Convert the trained PyTorch model to ONNX format for OpenVINO inference

To evaluate performance with OpenVINO, we can either convert the trained PyTorch model to an intermediate representation (IR) format or to an ONNX representation. In this notebook, we use the ONNX format. For more details on model optimization, refer to: https://docs.openvino.ai/latest/openvino_docs_MO_DG_Deep_Learning_Model_Optimizer_DevGuide.html

print('Converting the trained conve model to ONNX format')
torch.onnx.export(model, (torch.tensor(1), torch.tensor(1)),
                  fp32_onnx_path, verbose=False, opset_version=11, training=False)
Converting the trained conve model to ONNX format
/opt/home/k8sworker/cibuilds/ov-notebook/OVNotebookOps-231/.workspace/scm/ov-notebook/.venv/lib/python3.8/site-packages/torch/onnx/utils.py:305: UserWarning: It is recommended that constant folding be turned off ('do_constant_folding=False') when exporting the model in training-amenable mode, i.e. with 'training=TrainingMode.TRAIN' or 'training=TrainingMode.PRESERVE' (when model is in training mode). Otherwise, some learnable model parameters may not translate correctly in the exported ONNX model because constant folding mutates model parameters. Please consider turning off constant folding or setting the training=TrainingMode.EVAL.
  warnings.warn("It is recommended that constant folding be turned off ('do_constant_folding=False') "

Evaluate the model performance with OpenVINO

Now, we evaluate the model performance with the OpenVINO framework. In order to do so, we make three main API calls: 1. Initialize the Inference engine with Core() 2. Load the model with read_model() 3. Compile the model with compile_model()

The model can then be inferred on using by using the create_infer_request() API call.

ie = Core()
ir_net = ie.read_model(model=fp32_onnx_path)
compiled_model = ie.compile_model(model=ir_net)
input_layer_source = compiled_model.input('input.1')
input_layer_relation = compiled_model.input('input.2')
output_layer = compiled_model.output(0)

ov_acc = 0.0
ov_inf_times = []
for i in range(num_test_samples):
    test_sample = triples_list[i]
    source, relation, target = test_sample
    model_inputs = {input_layer_source: source, input_layer_relation: relation}
    start_time = time.time()
    result = compiled_model(model_inputs)[output_layer]
    end_time = time.time()
    ov_inf_times.append(end_time - start_time)
    top_k_idxs = list(np.argpartition(result[0], -TOP_K)[-TOP_K:])

    gt = np.array(sorted(t))
    pred = np.array(sorted(top_k_idxs))
    ov_acc += accuracy_score(gt, pred)

avg_ov_time = np.mean(ov_inf_times) * 1000
print(f'Average time taken for inference: {avg_ov_time} ms')
print(f'Mean accuracy of the model on the test dataset: {ov_acc/num_test_samples}')
Average time taken for inference: 1.7911593119303386 ms
Mean accuracy of the model on the test dataset: 0.10416666666666667

Determine the platform specific speedup obtained through OpenVINO graph optimizations

print(f'Speedup with OpenVINO optimizations: {round(float(avg_pt_time)/float(avg_ov_time),2)} X')
Speedup with OpenVINO optimizations: 0.35 X

Benchmark the converted OpenVINO model using benchmark app

The OpenVINO toolkit provides a benchmarking application to gauge the platform specific runtime performance that can be obtained under optimal configuration parameters for a given model. For more details refer to: https://docs.openvino.ai/latest/openvino_inference_engine_tools_benchmark_tool_README.html

Here, we use the benchmark application to obtain performance estimates under optimal configuration for the knowledge graph model inference. We obtain the average (AVG), minimum (MIN) as well as maximum (MAX) latency as well as the throughput performance (in samples/s) observed while running the benchmark application. The platform specific optimal configuration parameters determined by the benchmarking app for OpenVINO inference can also be obtained by looking at the benchmark app results.

print('Benchmark OpenVINO model using the benchmark app')
! benchmark_app -m "$fp32_onnx_path" -d CPU -api async -t 10 -shape "input.1[1],input.2[1]"
Benchmark OpenVINO model using the benchmark app
[Step 1/11] Parsing and validating input arguments
[ WARNING ]  -nstreams default value is determined automatically for a device. Although the automatic selection usually provides a reasonable performance, but it still may be non-optimal for some cases, for more information look at README.
[Step 2/11] Loading OpenVINO
[ WARNING ] PerformanceMode was not explicitly specified in command line. Device CPU performance hint will be set to THROUGHPUT.
[ INFO ] OpenVINO:
         API version............. 2022.1.0-7019-cdb9bec7210-releases/2022/1
[ INFO ] Device info
         CPU
         openvino_intel_cpu_plugin version 2022.1
         Build................... 2022.1.0-7019-cdb9bec7210-releases/2022/1

[Step 3/11] Setting device configuration
[ WARNING ] -nstreams default value is determined automatically for CPU device. Although the automatic selection usually provides a reasonable performance, but it still may be non-optimal for some cases, for more information look at README.
[Step 4/11] Reading network files
[ INFO ] Read model took 27.51 ms
[Step 5/11] Resizing network to match image sizes and given batch
[ INFO ] Reshaping model: 'input.1': {1}, 'input.2': {1}
[ INFO ] Reshape model took 0.82 ms
[ INFO ] Network batch size: 1
[Step 6/11] Configuring input of the model
[ INFO ] Model input 'input.1' precision i64, dimensions ([...]): 1
[ INFO ] Model input 'input.2' precision i64, dimensions ([...]): 1
[ INFO ] Model output '51' precision f32, dimensions ([...]): 1 271
[Step 7/11] Loading the model to the device
[ INFO ] Compile model took 50.77 ms
[Step 8/11] Querying optimal runtime parameters
[ INFO ] DEVICE: CPU
[ INFO ]   AVAILABLE_DEVICES  , ['']
[ INFO ]   RANGE_FOR_ASYNC_INFER_REQUESTS  , (1, 1, 1)
[ INFO ]   RANGE_FOR_STREAMS  , (1, 24)
[ INFO ]   FULL_DEVICE_NAME  , Intel(R) Core(TM) i9-10920X CPU @ 3.50GHz
[ INFO ]   OPTIMIZATION_CAPABILITIES  , ['WINOGRAD', 'FP32', 'FP16', 'INT8', 'BIN', 'EXPORT_IMPORT']
[ INFO ]   CACHE_DIR  ,
[ INFO ]   NUM_STREAMS  , 6
[ INFO ]   INFERENCE_NUM_THREADS  , 0
[ INFO ]   PERF_COUNT  , False
[ INFO ]   PERFORMANCE_HINT_NUM_REQUESTS  , 0
[Step 9/11] Creating infer requests and preparing input data
[ INFO ] Create 6 infer requests took 1.78 ms
[ WARNING ] No input files were given for input 'input.1'!. This input will be filled with random values!
[ WARNING ] No input files were given for input 'input.2'!. This input will be filled with random values!
[ INFO ] Fill input 'input.1' with random values
[ INFO ] Fill input 'input.2' with random values
[Step 10/11] Measuring performance (Start inference asynchronously, 6 inference requests using 6 streams for CPU, inference only: True, limits: 10000 ms duration)
[ INFO ] Benchmarking in inference only mode (inputs filling are not included in measurement loop).
[ INFO ] First inference took 4.07 ms
[Step 11/11] Dumping statistics report
Count:          84108 iterations
Duration:       10000.84 ms
Latency:
    Median:     0.62 ms
    AVG:        0.64 ms
    MIN:        0.41 ms
    MAX:        2.45 ms
Throughput: 8410.09 FPS

Conclusions

In this notebook, we converted the trained PyTorch knowledge graph embeddings model to the OpenVINO format. We confirmed that there were no accuracy differences post conversion. We also performed a sample evaluation on the knowledge graph! We then determined the platform specific speedup in runtime performance that can be obtained through OpenVINO graph optimizations. To learn more about the OpenVINO performance optimizations, refer to: https://docs.openvino.ai/latest/openvino_docs_optimization_guide_dldt_optimization_guide.html

References

  1. Convolutional 2D Knowledge Graph Embeddings, Tim Dettmers et al. (https://arxiv.org/abs/1707.01476)

  2. Model implementation: https://github.com/TimDettmers/ConvE

The ConvE model implementation used in this notebook licensed under the MIT License. The license is displayed below: MIT License

Copyright (c) 2017 Tim Dettmers

Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the “Software”), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED “AS IS”, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.