openvino-images-custom-ja.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "e296d68c-274d-40af-9f8d-5116d0587e96",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "# 分散モデル推論を用いた OpenVINO の使用\n",
    "このノートブックは、[torchvision.models](https://pytorch.org/docs/stable/torchvision/models.html#torchvision.models.resnet50)からResNet-50モデルを使用して、OpenVINOで分散モデル推論を行う方法をデモンストレーションします。入力データとして画像ファイルを使用します。\n",
    "\n",
    "このガイドは以下のセクションで構成されています：\n",
    "\n",
    "* **推論のための訓練済みモデルの準備。**\n",
    "* **Spark DataFrames に databricks-dataset からデータをロード。**\n",
    "* **Pandas UDF を用いたモデル推論の実行。**\n",
    "\n",
    "**注意:**\n",
    "* CPU対応のApache Sparkクラスタ上でノートブックを実行するには、変数`cuda = False`に変更してください。\n",
    "* GPU対応のApache Sparkクラスタ上でノートブックを実行するには、変数`cuda = True`に変更してください。\n",
    "* DBR 13.3 LTS ML, Standard_D16s_v5（ドライバー１台、ワーカー４台）"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "1cd78bf6-73f0-49a0-a9e7-75adcf084736",
     "showTitle": false,
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "%pip install openvino"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "524146a9-5755-4bc1-b808-0bc4612c4bb4",
     "showTitle": false,
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "dbutils.library.restartPython()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "0a7e1ca1-954d-47c2-a2c8-00d76c74447e",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "### 必要なライブラリをインポート"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "fb452933-ad97-4ee0-a9e6-c2296bc7aeb9",
     "showTitle": false,
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import shutil\n",
    "import uuid\n",
    "from typing import Iterator, Tuple\n",
    "\n",
    "import pandas as pd\n",
    "\n",
    "import torch\n",
    "from torch.utils.data import Dataset\n",
    "from torchvision import datasets, models, transforms\n",
    "from torchvision.datasets.folder import default_loader  # private API\n",
    "\n",
    "from pyspark.sql.functions import col, pandas_udf\n",
    "from pyspark.sql.types import ArrayType, FloatType, IntegerType\n",
    "\n",
    "import numpy as np\n",
    "\n",
    "import openvino as ov"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "7c19dd57-86d0-40bd-952c-d728f7c4af3f",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "### GPUを使用する場合はTrueに設定"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "0d892cc9-b206-4129-b6ec-6e2ce366ea46",
     "showTitle": false,
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "cuda = False\n",
    "\n",
    "use_cuda = cuda and torch.cuda.is_available()\n",
    "device = torch.device(\"cuda\" if use_cuda else \"cpu\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "ff0aaf6f-6616-497b-839f-d68f8ed1b0b2",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "### 学習済みモデルを推論用に準備する"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "3850a14e-ea78-47b2-9eb7-da2340eb8c03",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "ドライバーノードにResNet50をロードし、OpenVINOのフォーマットに変換してから、ワークスペースに保存する。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "56cad39f-5b88-4d38-aba4-e4553c7f8567",
     "showTitle": false,
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "# Save openvino.runtime.Model object on disk\n",
    "model = models.resnet50(weights=models.ResNet50_Weights.DEFAULT)\n",
    "model.eval()\n",
    "\n",
    "# Create OpenVINO Core object instance\n",
    "core = ov.Core()\n",
    "\n",
    "# Convert model to openvino.runtime.Model object\n",
    "ov_model = ov.convert_model(model)\n",
    "\n",
    "OV_MODEL_PATH = '/Workspace/Users/hiroshi.ouchiyama@databricks.com/ov_model/ov_resnet50_dynamic.xml'\n",
    "ov.save_model(ov_model, OV_MODEL_PATH)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "814563cc-9266-4afb-b44b-a6d7b3665895",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "モデルインスタンス作成用の関数を定義"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "d802c948-bcf6-4bee-b82c-396a01691a02",
     "showTitle": false,
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "def get_model_for_eval():\n",
    "  # Create OpenVINO Core object instance\n",
    "  core = ov.Core()\n",
    "  ov_model = core.read_model(model=OV_MODEL_PATH)\n",
    "\n",
    "  # Load OpenVINO model on device\n",
    "  compiled_model = core.compile_model(ov_model, 'CPU')\n",
    "\n",
    "  return compiled_model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "736bfbe7-2e23-441a-9546-d5db8ee97d95",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "### Spark DataFramesへのdatabricks-datasetからのデータのロード\n",
    "例として、TensorFlowチームによる[flowers dataset](https://www.tensorflow.org/datasets/catalog/tf_flowers)を使用します。これには、クラスごとに1つずつ、計5つのサブディレクトリの下に保存された花の写真が含まれています。これは、簡単にアクセスできるようにDatabricks Datasetsの`dbfs:/databricks-datasets/flower_photos`の下にホストされています。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "b2fd1900-5005-4825-9a05-d0069bc1c2a1",
     "showTitle": false,
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "dataset_dir = \"/dbfs/databricks-datasets/flower_photos/\"\n",
    "output_file_path = \"/tmp/predictions\"\n",
    "files = [os.path.join(dp, f) for dp, dn, filenames in os.walk(dataset_dir) for f in filenames if os.path.splitext(f)[1] == '.jpg']\n",
    "print(f'画像ファイルの総数は　{len(files)}　枚です。')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "54c9bad8-9355-45fc-94ec-2d3ff9cd170c",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "### シングルノード、または、ドライバーノード上でのみ推論させる"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "f989c3fb-05b1-4daa-b278-2f7237cbb817",
     "showTitle": false,
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "transform = transforms.Compose([\n",
    "  transforms.Resize(224),\n",
    "  transforms.CenterCrop(224),\n",
    "  transforms.ToTensor(),\n",
    "  transforms.Normalize(mean=[0.485, 0.456, 0.406],\n",
    "                      std=[0.229, 0.224, 0.225])\n",
    "])\n",
    "\n",
    "model = get_model_for_eval()\n",
    "\n",
    "predictions = []\n",
    "for image_path in files:\n",
    "  image = default_loader(image_path)\n",
    "  image = transform(image)\n",
    "  batch = image.unsqueeze(0)\n",
    "\n",
    "  prediction = model(batch)[0]\n",
    "  class_id = prediction.argmax(axis=1)\n",
    "  score = prediction[np.arange(prediction.shape[0]), class_id]\n",
    "  predictions.append((class_id, score))\n",
    "\n",
    "print(list(zip(files, predictions)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "970e7a9f-c9f0-4e27-9214-35473b568a77",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "### ここからようやく本題\n",
    "### 次にワーカーのノードを利用して分散モデル推論を実施する\n",
    "### Databricksでは分散モデル推論にはPandas UDF（別名：Vectorized UDF）を使うことが推奨されています。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "6f5d6ca7-d1de-49dd-a869-b304b6090531",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "分散処理のために、画像パスのDataFrameを作成する。\n",
    "\n",
    "再パーティショニングの際のパーティション数はワーカーノード数とおなじ、または最小限の倍数にすべし。このサンプルでは４つのワーカーノード使用しているので４に設定。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "189a021a-9c6a-42b2-a08f-49d2d0acf49e",
     "showTitle": false,
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "files_df = spark.createDataFrame(\n",
    "  map(lambda path: (path,), files), [\"path\"]\n",
    ").repartition(4)  # number of partitions should be a small multiple of total number of nodes　\n",
    "\n",
    "display(files_df.limit(10))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "17330808-e2be-489d-9a23-2c08c41ed03f",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "各画像パスのDataFrameが４つのパーティションにほぼ等分されているのが確認できます。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "30a94680-5796-4fcd-9373-8c02f2ca10a8",
     "showTitle": false,
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "from pyspark.sql.functions import spark_partition_id\n",
    "\n",
    "display(files_df.withColumn('partition', spark_partition_id()).groupBy('partition').count().orderBy('partition'))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "febaf184-ecd0-4122-a08e-8bbf07f36e75",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "### Pandas UDFを作成します。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "3eb871bc-1886-4edc-af55-adc07c0f977a",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "Pandas UDFはメモリ内のデータ形式にApache Arrowを使っているので、明示的にArrowの使用をOnにします。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "62696bc0-faff-4aac-82a1-c691ce4452b2",
     "showTitle": false,
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "spark.conf.set(\"spark.sql.execution.arrow.enabled\", \"true\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "e0a478e9-c0c8-431e-930a-ad8f584a4e0f",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "かつ、Pandas UDFに一度に渡すレコード数（バッチ）の最大値を`spark.sql.execution.arrow.maxRecordsPerBatch`に指定します。今回は５１２を設定しています。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "7b6331ae-7697-4010-a074-6196c63815d3",
     "showTitle": false,
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "spark.conf.set(\"spark.sql.execution.arrow.maxRecordsPerBatch\", \"５１２\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "0e5de115-6ceb-4218-98f6-d0bcf57ee430",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "カスタム PyTorch データセットクラスを作成します。\n",
    "\n",
    "以下のドキュメントの通りですが、パフォーマンスチューニングのヒントとして、PyTorchであれば、データロード用に`torch.utils.data.DataLoader`の使用が推奨されているので、そのお作法に則ります。\n",
    "- https://docs.databricks.com/en/machine-learning/model-inference/dl-model-inference.html\n",
    "- https://docs.databricks.com/en/machine-learning/model-inference/model-inference-performance.html"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "fe87bf01-38f3-4e3f-86ab-dcd03944ac3d",
     "showTitle": false,
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "class ImageDataset(Dataset):\n",
    "  def __init__(self, paths, transform=None):\n",
    "    self.paths = paths\n",
    "    self.transform = transform\n",
    "  def __len__(self):\n",
    "    return len(self.paths)\n",
    "  def __getitem__(self, index):\n",
    "    image = default_loader(self.paths[index])\n",
    "    if self.transform is not None:\n",
    "      image = self.transform(image)\n",
    "    return image"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "fe906323-4ce6-475c-a358-bc61f49c47e7",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "モデル推論のための関数を定義する。\n",
    "\n",
    "今回のサンプル画像データは全部で3670枚で、ワーカーノードの数と同じ４つのパーティションに分割しているので、各ワーカーノードが９１０〜９２０枚ほどの画像を処理します。その中から512枚の画像（正確には画像パス）を取り出してきて、Pandas UDFにpadas.Seriesデータとして入力します。Pandas UDF内では、その５１２個の画像パスからバッチサイズごとに画像パスを取り出し、当該画像ファイルをロードしてTensor化して、それをモデルで推論します。推論結果から欲しい情報を取り出した上で、それをPandas.Seriesとしてパックして、返します。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "cd093ebe-69da-472d-9072-f4028e9e6b09",
     "showTitle": false,
     "title": ""
    }
   },
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "07e14852-d5d5-4cb5-9abd-45b5bcbeb3a6",
     "showTitle": false,
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "@pandas_udf(ArrayType(FloatType()))\n",
    "def predict_batch_udf(paths: pd.Series) -> pd.Series:\n",
    "  transform = transforms.Compose([\n",
    "    transforms.Resize(224),\n",
    "    transforms.CenterCrop(224),\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize(mean=[0.485, 0.456, 0.406],\n",
    "                       std=[0.229, 0.224, 0.225])\n",
    "  ])\n",
    "\n",
    "  images = ImageDataset(paths, transform=transform)\n",
    "  loader = torch.utils.data.DataLoader(images, batch_size=1, num_workers=8)\n",
    "  model = get_model_for_eval()\n",
    "\n",
    "  all_predictions = []\n",
    "  with torch.no_grad():\n",
    "    for batch in loader:\n",
    "      predictions = model(batch)[0]\n",
    "      class_id = predictions.argmax(axis=1)\n",
    "      score = predictions[np.arange(predictions.shape[0]), class_id]\n",
    "\n",
    "      for result in np.stack((class_id, score), axis=1):\n",
    "        all_predictions.append(result)\n",
    "\n",
    "  return pd.Series(all_predictions)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "55fe4dd1-079e-4c98-bb64-7566402c7418",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "モデル推論を実行し、結果をDisplayします。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "74c071d1-3369-498c-a141-1aa3ef1f9dc8",
     "showTitle": false,
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "predictions_df = files_df.withColumn('prediction', predict_batch_udf(col('path')))\n",
    "display(predictions_df)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "4295f693-a5b0-4785-a7a8-f4904a6b0712",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "以上、Databricks上で分散モデル推論を実施する際に、Pandas UDFを使用する方法をご紹介しました。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "c3383862-db57-4384-b4b5-93a25a803222",
     "showTitle": false,
     "title": ""
    }
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "application/vnd.databricks.v1+notebook": {
   "dashboards": [],
   "language": "python",
   "notebookMetadata": {
    "pythonIndentUnit": 2
   },
   "notebookName": "openvino-images-custom-ja",
   "widgets": {}
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}