notebooks/219-knowledge-graphs-conve/219-knowledge-graphs-conve.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# OpenVINO optimizations for Knowledge graphs\n",
    "\n",
    "The goal of this notebook is to showcase performance optimizations for the ConvE knowledge graph embeddings model using the Intel® Distribution of OpenVINO™ Toolkit. <br>\n",
    "The optimizations process contains the following steps:\n",
    "1. Export the trained model to a format suitable for OpenVINO optimizations and inference\n",
    "2. Report the inference performance speedup obtained with the optimized OpenVINO model\n",
    "\n",
    "The ConvE model is an implementation of the paper - \"Convolutional 2D Knowledge Graph Embeddings\" (https://arxiv.org/abs/1707.01476). <br>\n",
    "The sample dataset can be downloaded from: https://github.com/TimDettmers/ConvE/tree/master/countries/countries_S1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Windows specific settings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# On Windows, add the directory that contains cl.exe to the PATH\n",
    "# to enable PyTorch to find the required C++ tools.\n",
    "# This code assumes that Visual Studio 2019 is installed in the default directory.\n",
    "# If you have a different C++ compiler, please add the correct path\n",
    "# to os.environ[\"PATH\"] directly.\n",
    "# Note that the C++ Redistributable is not enough to run this notebook.\n",
    "\n",
    "# Adding the path to os.environ[\"LIB\"] is not always required\n",
    "# - it depends on the system's configuration\n",
    "\n",
    "import sys\n",
    "\n",
    "if sys.platform == \"win32\":\n",
    "    import distutils.command.build_ext\n",
    "    import os\n",
    "    from pathlib import Path\n",
    "\n",
    "    VS_INSTALL_DIR = r\"C:/Program Files (x86)/Microsoft Visual Studio\"\n",
    "    cl_paths = sorted(list(Path(VS_INSTALL_DIR).glob(\"**/Hostx86/x64/cl.exe\")))\n",
    "    if len(cl_paths) == 0:\n",
    "        raise ValueError(\n",
    "            \"Cannot find Visual Studio. This notebook requires a C++ compiler. If you installed \"\n",
    "            \"a C++ compiler, please add the directory that contains\"\n",
    "            \"cl.exe to `os.environ['PATH']`.\"\n",
    "        )\n",
    "    else:\n",
    "        # If multiple versions of MSVC are installed, get the most recent version\n",
    "        cl_path = cl_paths[-1]\n",
    "        vs_dir = str(cl_path.parent)\n",
    "        os.environ[\"PATH\"] += f\"{os.pathsep}{vs_dir}\"\n",
    "        # Code for finding the library dirs from\n",
    "        # https://stackoverflow.com/questions/47423246/get-pythons-lib-path\n",
    "        d = distutils.core.Distribution()\n",
    "        b = distutils.command.build_ext.build_ext(d)\n",
    "        b.finalize_options()\n",
    "        os.environ[\"LIB\"] = os.pathsep.join(b.library_dirs)\n",
    "        print(f\"Added {vs_dir} to PATH\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Import the packages needed for successful execution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import time\n",
    "import json\n",
    "\n",
    "import torch\n",
    "from torch.nn import functional as F, Parameter\n",
    "from torch.nn.init import xavier_normal_\n",
    "\n",
    "from pathlib import Path\n",
    "\n",
    "from sklearn.metrics import accuracy_score\n",
    "\n",
    "from openvino.runtime import Core"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Settings: Including path to the serialized model files and input data files\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "print(f\"Using {device} device\")\n",
    "\n",
    "# Path to the trained model\n",
    "modelpath = Path('models/conve.pt')\n",
    "\n",
    "# Path to the file containing the entities and entity IDs\n",
    "entdatapath = Path('../data/text/countries_S1/kg_training_entids.txt')\n",
    "\n",
    "# Path to the file containing the relations and relation IDs\n",
    "reldatapath = Path('../data/text/countries_S1/kg_training_relids.txt')\n",
    "\n",
    "# Path to the test data file\n",
    "testdatapath = Path('../data/json/countries_S1/e1rel_to_e2_ranking_test.json')\n",
    "\n",
    "# Entity and relation embedding dimensions\n",
    "EMB_DIM = 300\n",
    "\n",
    "# Top K vals to consider from the predictions\n",
    "TOP_K = 2\n",
    "\n",
    "# Required for OpenVINO conversion\n",
    "output_dir = Path(\"models\")\n",
    "base_model_name = \"conve\"\n",
    "\n",
    "output_dir.mkdir(exist_ok=True)\n",
    "\n",
    "# Paths where PyTorch, ONNX and OpenVINO IR models will be stored\n",
    "fp32_onnx_path = Path(output_dir / (base_model_name + \"_fp32\")).with_suffix(\".onnx\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Defining the ConvE model class"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Model implementation reference: https://github.com/TimDettmers/ConvE\n",
    "class ConvE(torch.nn.Module):\n",
    "    def __init__(self, num_entities, num_relations, emb_dim):\n",
    "        super(ConvE, self).__init__()\n",
    "        # Embedding tables for entity and relations with num_uniq_ent in y-dim, emb_dim in x-dim\n",
    "        self.emb_e = torch.nn.Embedding(num_entities, emb_dim, padding_idx=0)\n",
    "        self.ent_weights_matrix = torch.ones([num_entities, emb_dim], dtype=torch.float64)\n",
    "        self.emb_rel = torch.nn.Embedding(num_relations, emb_dim, padding_idx=0)\n",
    "        self.ne = num_entities\n",
    "        self.nr = num_relations\n",
    "        self.inp_drop = torch.nn.Dropout(0.2)\n",
    "        self.hidden_drop = torch.nn.Dropout(0.3)\n",
    "        self.feature_map_drop = torch.nn.Dropout2d(0.2)\n",
    "        self.loss = torch.nn.BCELoss()\n",
    "        self.conv1 = torch.nn.Conv2d(1, 32, (3, 3), 1, 0, bias=True)\n",
    "        self.bn0 = torch.nn.BatchNorm2d(1)\n",
    "        self.bn1 = torch.nn.BatchNorm2d(32)\n",
    "        self.ln0 = torch.nn.LayerNorm(emb_dim)\n",
    "        self.register_parameter('b', Parameter(torch.zeros(num_entities)))\n",
    "        self.fc = torch.nn.Linear(16128, emb_dim)\n",
    "\n",
    "    def init(self):\n",
    "        \"\"\" Initializes the model \"\"\"\n",
    "        # Xavier initialization\n",
    "        xavier_normal_(self.emb_e.weight.data)\n",
    "        xavier_normal_(self.emb_rel.weight.data)\n",
    "\n",
    "    def forward(self, e1, rel):\n",
    "        \"\"\" Forward pass on the model.\n",
    "        :param e1: source entity\n",
    "        :param rel: relation between the source and target entities\n",
    "        Returns the model predictions for the target entities\n",
    "        \"\"\"\n",
    "        e1_embedded = self.emb_e(e1).view(-1, 1, 10, 30)\n",
    "        rel_embedded = self.emb_rel(rel).view(-1, 1, 10, 30)\n",
    "        stacked_inputs = torch.cat([e1_embedded, rel_embedded], 2)\n",
    "        stacked_inputs = self.bn0(stacked_inputs)\n",
    "        x = self.inp_drop(stacked_inputs)\n",
    "        x = self.conv1(x)\n",
    "        x = self.bn1(x)\n",
    "        x = F.relu(x)\n",
    "        x = self.feature_map_drop(x)\n",
    "        x = x.view(1, -1)\n",
    "        x = self.fc(x)\n",
    "        x = self.hidden_drop(x)\n",
    "        x = self.ln0(x)\n",
    "        x = F.relu(x)\n",
    "        x = torch.mm(x, self.emb_e.weight.transpose(1, 0))\n",
    "        x = self.hidden_drop(x)\n",
    "        x += self.b.expand_as(x)\n",
    "        pred = torch.nn.functional.softmax(x, dim=1)\n",
    "        return pred"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Defining the dataloader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DataLoader():\n",
    "    def __init__(self):\n",
    "        super(DataLoader, self).__init__()\n",
    "\n",
    "        self.ent_path = entdatapath\n",
    "        self.rel_path = reldatapath\n",
    "        self.test_file = testdatapath\n",
    "        self.entity_ids, self.ids2entities = self.load_data(data_path=self.ent_path)\n",
    "        self.rel_ids, self.ids2rel = self.load_data(data_path=self.rel_path)\n",
    "        self.test_triples_list = self.convert_triples(data_path=self.test_file)\n",
    "\n",
    "    def load_data(self, data_path):\n",
    "        \"\"\" Creates a dictionary of data items with corresponding ids \"\"\"\n",
    "        item_dict, ids_dict = {}, {}\n",
    "        fp = open(data_path, \"r\")\n",
    "        lines = fp.readlines()\n",
    "        for line in lines:\n",
    "            name, id = line.strip().split('\\t')\n",
    "            item_dict[name] = int(id)\n",
    "            ids_dict[int(id)] = name\n",
    "        fp.close()\n",
    "        return item_dict, ids_dict\n",
    "\n",
    "    def convert_triples(self, data_path):\n",
    "        \"\"\" Creates a triple of source entity, relation and target entities\"\"\"\n",
    "        triples_list = []\n",
    "        dp = open(data_path, \"r\")\n",
    "        lines = dp.readlines()\n",
    "        for line in lines:\n",
    "            item_dict = json.loads(line.strip())\n",
    "            h = item_dict['e1']\n",
    "            r = item_dict['rel']\n",
    "            t = item_dict['e2_multi1'].split('\\t')\n",
    "            hrt_list = []\n",
    "            hrt_list.append(self.entity_ids[h])\n",
    "            hrt_list.append(self.rel_ids[r])\n",
    "            t_ents = []\n",
    "            for t_idx in t:\n",
    "                t_ents.append(self.entity_ids[t_idx])\n",
    "            hrt_list.append(t_ents)\n",
    "            triples_list.append(hrt_list)\n",
    "        dp.close()\n",
    "        return triples_list"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Evaluate the trained ConvE model\n",
    "First, we will 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. <br>\n",
    "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."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = DataLoader()\n",
    "num_entities = len(data.entity_ids)\n",
    "num_relations = len(data.rel_ids)\n",
    "\n",
    "model = ConvE(num_entities=num_entities, num_relations=num_relations, emb_dim=EMB_DIM)\n",
    "model.load_state_dict(torch.load(modelpath))\n",
    "model.eval()\n",
    "\n",
    "pt_inf_times = []\n",
    "\n",
    "triples_list = data.test_triples_list\n",
    "num_test_samples = len(triples_list)\n",
    "pt_acc = 0.0\n",
    "for i in range(num_test_samples):\n",
    "    test_sample = triples_list[i]\n",
    "    h, r, t = test_sample\n",
    "    start_time = time.time()\n",
    "    logits = model.forward(e1=torch.tensor(h), rel=torch.tensor(r))\n",
    "    end_time = time.time()\n",
    "    pt_inf_times.append(end_time - start_time)\n",
    "    score, pred = torch.topk(logits, TOP_K, 1)\n",
    "\n",
    "    gt = np.array(sorted(t))\n",
    "    pred = np.array(sorted(pred[0].cpu().detach()))\n",
    "    pt_acc += accuracy_score(gt, pred)\n",
    "\n",
    "avg_pt_time = np.mean(pt_inf_times) * 1000\n",
    "print(f'Average time taken for inference: {avg_pt_time} ms')\n",
    "print(f'Mean accuracy of the model on the test dataset: {pt_acc/num_test_samples}')"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Prediction on the Knowledge graph.\n",
    "Here, we perform the entity prediction on the knowledge graph, as a sample evaluation task. <br>\n",
    "We pass the source entity 'san_marino' and relation 'locatedIn' to the knowledge graph and obtain the target entity predictions. <br>\n",
    "Expected predictions are target entities that form a factual triple with the entity and relation passed as inputs to the knowledge graph."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "entitynames_dict = data.ids2entities\n",
    "\n",
    "ent = 'san_marino'\n",
    "rel = 'locatedin'\n",
    "\n",
    "h_idx = data.entity_ids[ent]\n",
    "r_idx = data.rel_ids[rel]\n",
    "\n",
    "logits = model.forward(torch.tensor(h_idx), torch.tensor(r_idx))\n",
    "score, pred = torch.topk(logits, TOP_K, 1)\n",
    "\n",
    "for j, id in enumerate(pred[0].cpu().detach().numpy()):\n",
    "    pred_entity = entitynames_dict[id]\n",
    "    print(f'Source Entity: {ent}, Relation: {rel}, Target entity prediction: {pred_entity}')"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Convert the trained PyTorch model to ONNX format for OpenVINO inference\n",
    "To evaluate performance with OpenVINO, we can either convert the trained PyTorch model to an intermediate representation (IR) format or to an ONNX representation. <br>\n",
    "This notebook uses 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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('Converting the trained conve model to ONNX format')\n",
    "torch.onnx.export(model, (torch.tensor(1), torch.tensor(1)),\n",
    "                  fp32_onnx_path, input_names=['input.1', 'input.2'], verbose=False, opset_version=11)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Evaluate the model performance with OpenVINO\n",
    "Now, we evaluate the model performance with the OpenVINO framework. In order to do so, make three main API calls:\n",
    "1. Initialize the Inference engine with `Core()`\n",
    "2. Load the model with `read_model()`\n",
    "3. Compile the model with `compile_model()`\n",
    "\n",
    "Then, the model can be inferred on by using the `create_infer_request()` API call."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ie = Core()\n",
    "ir_net = ie.read_model(model=fp32_onnx_path)\n",
    "compiled_model = ie.compile_model(model=ir_net)\n",
    "input_layer_source = compiled_model.input('input.1')\n",
    "input_layer_relation = compiled_model.input('input.2')\n",
    "output_layer = compiled_model.output(0)\n",
    "\n",
    "ov_acc = 0.0\n",
    "ov_inf_times = []\n",
    "for i in range(num_test_samples):\n",
    "    test_sample = triples_list[i]\n",
    "    source, relation, target = test_sample\n",
    "    model_inputs = {input_layer_source: np.int64(source), input_layer_relation: np.int64(relation)}\n",
    "    start_time = time.time()\n",
    "    result = compiled_model(model_inputs)[output_layer]\n",
    "    end_time = time.time()\n",
    "    ov_inf_times.append(end_time - start_time)\n",
    "    top_k_idxs = list(np.argpartition(result[0], -TOP_K)[-TOP_K:])\n",
    "\n",
    "    gt = np.array(sorted(t))\n",
    "    pred = np.array(sorted(top_k_idxs))\n",
    "    ov_acc += accuracy_score(gt, pred)\n",
    "\n",
    "avg_ov_time = np.mean(ov_inf_times) * 1000\n",
    "print(f'Average time taken for inference: {avg_ov_time} ms')\n",
    "print(f'Mean accuracy of the model on the test dataset: {ov_acc/num_test_samples}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Determine the platform specific speedup obtained through OpenVINO graph optimizations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f'Speedup with OpenVINO optimizations: {round(float(avg_pt_time)/float(avg_ov_time),2)} X')"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Benchmark the converted OpenVINO model using benchmark app\n",
    "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\n",
    "\n",
    "Here, we use the benchmark application to obtain performance estimates under optimal configuration for the knowledge graph model inference. <br> \n",
    "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. <br> \n",
    "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. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('Benchmark OpenVINO model using the benchmark app')\n",
    "! benchmark_app -m \"$fp32_onnx_path\" -d CPU -api async -t 10 -shape \"input.1[1],input.2[1]\" "
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Conclusions\n",
    "In this notebook, we convert the trained PyTorch knowledge graph embeddings model to the OpenVINO format. <br>\n",
    "We confirm that there are no accuracy differences post conversion. We also perform a sample evaluation on the knowledge graph.<br> \n",
    "Then, we determine the platform specific speedup in runtime performance that can be obtained through OpenVINO graph optimizations. <br>\n",
    "To learn more about the OpenVINO performance optimizations, refer to: https://docs.openvino.ai/latest/openvino_docs_optimization_guide_dldt_optimization_guide.html "
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### References\n",
    "1. Convolutional 2D Knowledge Graph Embeddings, Tim Dettmers et al. (https://arxiv.org/abs/1707.01476)\n",
    "2. Model implementation: https://github.com/TimDettmers/ConvE \n",
    "\n",
    "The ConvE model implementation used in this notebook is licensed under the MIT License. The license is displayed below: <br>\n",
    "MIT License\n",
    "\n",
    "Copyright (c) 2017 Tim Dettmers\n",
    "\n",
    "Permission is hereby granted, free of charge, to any person obtaining a copy<br>\n",
    "of this software and associated documentation files (the \"Software\"), to deal<br>\n",
    "in the Software without restriction, including without limitation the rights<br>\n",
    "to use, copy, modify, merge, publish, distribute, sublicense, and/or sell<br>\n",
    "copies of the Software, and to permit persons to whom the Software is<br>\n",
    "furnished to do so, subject to the following conditions:\n",
    "\n",
    "The above copyright notice and this permission notice shall be included in all<br>\n",
    "copies or substantial portions of the Software.\n",
    "\n",
    "THE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR<br>\n",
    "IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,<br>\n",
    "FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE<br>\n",
    "AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER<br>\n",
    "LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,<br>\n",
    "OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE<br>\n",
    "SOFTWARE."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.13 (main, Aug 25 2022, 23:51:50) [MSC v.1916 64 bit (AMD64)]"
  },
  "vscode": {
   "interpreter": {
    "hash": "e0404472fd7b5b63117a9fa5c50283296e2708c2449c6090d2cdf8903f95897f"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}