notebooks/ONNX - Inference on Spark.ipynb

{
    "cells": [
        {
            "cell_type": "markdown",
            "source": [
                "## ONNX Inference on Spark\r\n",
                "\r\n",
                "In this example, we will train a LightGBM model, convert the model to ONNX format and use the converted model to infer some testing data on Spark.\r\n",
                "\r\n",
                "Python dependencies:\r\n",
                "\r\n",
                "- onnxmltools==1.7.0\r\n",
                "- lightgbm==3.2.1\r\n"
            ],
            "metadata": {}
        },
        {
            "cell_type": "markdown",
            "source": [
                "Load training data"
            ],
            "metadata": {}
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "source": [
                "df = spark.read.format(\"csv\")\\\r\n",
                "  .option(\"header\", True)\\\r\n",
                "  .option(\"inferSchema\", True)\\\r\n",
                "  .load(\"wasbs://publicwasb@mmlspark.blob.core.windows.net/company_bankruptcy_prediction_data.csv\")\r\n",
                "\r\n",
                "display(df)"
            ],
            "outputs": [],
            "metadata": {}
        },
        {
            "cell_type": "markdown",
            "source": [
                "Use LightGBM to train a model"
            ],
            "metadata": {}
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "source": [
                "from pyspark.ml.feature import VectorAssembler\r\n",
                "from mmlspark.lightgbm import LightGBMClassifier\r\n",
                "\r\n",
                "feature_cols = df.columns[1:]\r\n",
                "featurizer = VectorAssembler(\r\n",
                "    inputCols=feature_cols,\r\n",
                "    outputCol='features'\r\n",
                ")\r\n",
                "\r\n",
                "train_data = featurizer.transform(df)['Bankrupt?', 'features']\r\n",
                "\r\n",
                "model = (\r\n",
                "  LightGBMClassifier(featuresCol=\"features\", labelCol=\"Bankrupt?\")\r\n",
                "  .setEarlyStoppingRound(300)\r\n",
                "  .setLambdaL1(0.5)\r\n",
                "  .setNumIterations(1000)\r\n",
                "  .setNumThreads(-1)\r\n",
                "  .setMaxDeltaStep(0.5)\r\n",
                "  .setNumLeaves(31)\r\n",
                "  .setMaxDepth(-1)\r\n",
                "  .setBaggingFraction(0.7)\r\n",
                "  .setFeatureFraction(0.7)\r\n",
                "  .setBaggingFreq(2)\r\n",
                "  .setObjective(\"binary\")\r\n",
                "  .setIsUnbalance(True)\r\n",
                "  .setMinSumHessianInLeaf(20)\r\n",
                "  .setMinGainToSplit(0.01)\r\n",
                ")\r\n",
                "\r\n",
                "model = model.fit(train_data)"
            ],
            "outputs": [],
            "metadata": {}
        },
        {
            "cell_type": "markdown",
            "source": [
                "Export the trained model to a LightGBM booster, convert it to ONNX format."
            ],
            "metadata": {}
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "source": [
                "import lightgbm as lgb\r\n",
                "from lightgbm import Booster, LGBMClassifier\r\n",
                "\r\n",
                "def convertModel(lgbm_model: LGBMClassifier or Booster, input_size: int) -> bytes:\r\n",
                "  from onnxmltools.convert import convert_lightgbm\r\n",
                "  from onnxconverter_common.data_types import FloatTensorType\r\n",
                "  initial_types = [(\"input\", FloatTensorType([-1, input_size]))]\r\n",
                "  onnx_model = convert_lightgbm(lgbm_model, initial_types=initial_types, target_opset=9)\r\n",
                "  return onnx_model.SerializeToString()\r\n",
                "\r\n",
                "booster_model_str = model.getLightGBMBooster().modelStr().get()\r\n",
                "booster = lgb.Booster(model_str=booster_model_str)\r\n",
                "model_payload_ml = convertModel(booster, len(df.columns) - 1)"
            ],
            "outputs": [],
            "metadata": {}
        },
        {
            "cell_type": "markdown",
            "source": [
                "Load the ONNX payload into an `ONNXModel`, and inspect the model inputs and outputs."
            ],
            "metadata": {}
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "source": [
                "from mmlspark.onnx import ONNXModel\r\n",
                "\r\n",
                "onnx_ml = ONNXModel().setModelPayload(model_payload_ml)\r\n",
                "\r\n",
                "print(\"Model inputs:\" + str(onnx_ml.getModelInputs()))\r\n",
                "print(\"Model outputs:\" + str(onnx_ml.getModelOutputs()))"
            ],
            "outputs": [],
            "metadata": {}
        },
        {
            "cell_type": "markdown",
            "source": [
                "Map the model input to the input dataframe's column name (FeedDict), and map the output dataframe's column names to the model outputs (FetchDict)."
            ],
            "metadata": {}
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "source": [
                "onnx_ml = (\r\n",
                "  onnx_ml\r\n",
                "    .setDeviceType(\"CPU\")\r\n",
                "    .setFeedDict({\"input\": \"features\"})\r\n",
                "    .setFetchDict({\"probability\": \"probabilities\", \"prediction\": \"label\"})\r\n",
                "    .setMiniBatchSize(5000)\r\n",
                ")"
            ],
            "outputs": [],
            "metadata": {}
        },
        {
            "cell_type": "markdown",
            "source": [
                "Create some testing data and transform the data through the ONNX model."
            ],
            "metadata": {}
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "source": [
                "from pyspark.ml.feature import VectorAssembler\r\n",
                "import pandas as pd\r\n",
                "import numpy as np\r\n",
                "\r\n",
                "n = 1000 * 1000\r\n",
                "m = 95\r\n",
                "test = np.random.rand(n, m)\r\n",
                "testPdf = pd.DataFrame(test)\r\n",
                "cols = list(map(str, testPdf.columns))\r\n",
                "testDf = spark.createDataFrame(testPdf)\r\n",
                "testDf = testDf.union(testDf).repartition(200)\r\n",
                "testDf = VectorAssembler().setInputCols(cols).setOutputCol(\"features\").transform(testDf).drop(*cols).cache()\r\n",
                "\r\n",
                "display(onnx_ml.transform(testDf))"
            ],
            "outputs": [],
            "metadata": {}
        }
    ],
    "metadata": {
        "language_info": {
            "name": "python"
        }
    },
    "nbformat": 4,
    "nbformat_minor": 2
}