tutorials/Tutorial_5_ONNX_inference.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "e75f5c6d",
   "metadata": {},
   "source": [
    "## This tutorial shows how to convert your model to ONNX and use for CPU inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "8f07b0d0",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "# Optional: set the device to run\n",
    "os.environ[\"CUDA_DEVICE_ORDER\"] = \"PCI_BUS_ID\"\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0\"\n",
    "\n",
    "os.makedirs('../data', exist_ok=True)\n",
    "\n",
    "import numpy as np\n",
    "import joblib\n",
    "import onnxruntime\n",
    "\n",
    "from sklearn.datasets import make_regression\n",
    "\n",
    "from py_boost import GradientBoosting\n",
    "from py_boost import pb_to_onnx, ONNXPredictor"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3aa67d64",
   "metadata": {},
   "source": [
    "### Generate dummy multilabel task and train the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "12c1d5ec",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[15:37:33] Stdout logging level is INFO.\n",
      "[15:37:33] GDBT train starts. Max iter 100, early stopping rounds 200\n",
      "[15:37:34] Iter 0; \n",
      "[15:37:37] Iter 99; \n",
      "CPU times: user 15.1 s, sys: 1.85 s, total: 16.9 s\n",
      "Wall time: 10.1 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "X, y = make_regression(150000, 100, n_targets=5, random_state=42)\n",
    "# binarize\n",
    "y = (y > y.mean(axis=0)).astype(np.float32)\n",
    "\n",
    "model = GradientBoosting(\n",
    "    'bce', lr=0.01, verbose=100, \n",
    "    ntrees=100, es=200, \n",
    ")\n",
    "model.fit(X, y)\n",
    "pp = model.predict(X)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "49202ae7",
   "metadata": {},
   "source": [
    "### Convert the model to ONNX\n",
    "\n",
    "The simpliest way to convert is using `pb_to_onnx` function. Just pass the `py-boost` model and path to store parsed model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "21e10069",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 100/100 [00:00<00:00, 1723.04it/s]\n"
     ]
    }
   ],
   "source": [
    "pb_to_onnx(model, '../data/pb_model.onnx')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8d385bbc",
   "metadata": {},
   "source": [
    "Once the parsing is completed, you can run `onnxruntime` session for inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "7d7d4a0c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 5.59 s, sys: 131 ms, total: 5.72 s\n",
      "Wall time: 395 ms\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "array([[0.6264308 , 0.41568166, 0.5388822 , 0.4261355 , 0.57804173],\n",
       "       [0.59586126, 0.42369062, 0.56585   , 0.57584757, 0.5392887 ],\n",
       "       [0.72726965, 0.67056704, 0.49255225, 0.6711969 , 0.635281  ],\n",
       "       ...,\n",
       "       [0.5112887 , 0.38028964, 0.4761739 , 0.52265   , 0.4513791 ],\n",
       "       [0.67362005, 0.54282206, 0.62851644, 0.6090929 , 0.7003519 ],\n",
       "       [0.56341565, 0.52830017, 0.41594115, 0.43341845, 0.42639387]],\n",
       "      dtype=float32)"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "# start session\n",
    "sess = onnxruntime.InferenceSession(\n",
    "    '../data/pb_model.onnx', \n",
    "    providers=[\"CPUExecutionProvider\"]\n",
    ")\n",
    "\n",
    "# run inference\n",
    "preds = sess.run(['Y'], {'X': X.astype(np.float32, copy=False)})[0]\n",
    "preds = 1 / (1 + np.exp(-preds))\n",
    "preds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "50d94691",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "2.3841858e-07"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.abs(preds - pp).max()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "37a4c052",
   "metadata": {},
   "source": [
    "***Note*** : by default, parser only collect the trees and base score info. So, it knows nothing about the postprocessing function, for example `sigmoid` in this case. That's why we apply sigmoid after inference part. But we can pass one of built-in `ONNX` post transforms: 'NONE', 'SOFTMAX', 'LOGISTIC', 'SOFTMAX_ZERO', or 'PROBIT' to avoid this step. Probably it is going to be more efficient:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "9d32635e",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 100/100 [00:00<00:00, 1670.84it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 5.58 s, sys: 178 ms, total: 5.76 s\n",
      "Wall time: 583 ms\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "array([[0.62643087, 0.41568172, 0.5388822 , 0.42613554, 0.57804173],\n",
       "       [0.5958613 , 0.42369062, 0.56584996, 0.57584757, 0.5392887 ],\n",
       "       [0.72726965, 0.67056704, 0.49255228, 0.6711969 , 0.6352811 ],\n",
       "       ...,\n",
       "       [0.5112887 , 0.3802896 , 0.47617394, 0.52265   , 0.45137918],\n",
       "       [0.67362005, 0.54282206, 0.6285165 , 0.6090929 , 0.7003519 ],\n",
       "       [0.56341565, 0.5283001 , 0.41594112, 0.43341845, 0.42639393]],\n",
       "      dtype=float32)"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%%time\n",
    "pb_to_onnx(model, '../data/pb_model.onnx', post_transform='LOGISTIC') # pass built-in post transform\n",
    "\n",
    "# start session\n",
    "sess = onnxruntime.InferenceSession(\n",
    "    '../data/pb_model.onnx', \n",
    "    providers=[\"CPUExecutionProvider\"]\n",
    ")\n",
    "\n",
    "# run inference\n",
    "preds = sess.run(['Y'], {'X': X.astype(np.float32, copy=False)})[0]\n",
    "preds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "6efcaaa3",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "2.3841858e-07"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.abs(preds - pp).max()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bdcc1486",
   "metadata": {},
   "source": [
    "***Filter outputs*** . Another option is to convert just a part of outputs to `ONNX`, for the case when we need only particular outputs for inference. For example, we want to keep only 0 and 2 outputs for inference and we don't want to compute the parts of model that doesn't affect the result:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "080f1139",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 100/100 [00:00<00:00, 2039.98it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 5.31 s, sys: 178 ms, total: 5.48 s\n",
      "Wall time: 528 ms\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "array([[0.62643087, 0.5388822 ],\n",
       "       [0.5958613 , 0.56584996],\n",
       "       [0.72726965, 0.49255228],\n",
       "       ...,\n",
       "       [0.5112887 , 0.47617394],\n",
       "       [0.67362005, 0.6285165 ],\n",
       "       [0.56341565, 0.41594112]], dtype=float32)"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%%time\n",
    "pb_to_onnx(model, '../data/pb_model.onnx', fltr=[0, 2], post_transform='LOGISTIC') # pass array to filter outputs\n",
    "\n",
    "# start session\n",
    "sess = onnxruntime.InferenceSession(\n",
    "    '../data/pb_model.onnx', \n",
    "    providers=[\"CPUExecutionProvider\"]\n",
    ")\n",
    "\n",
    "# run inference\n",
    "preds = sess.run(['Y'], {'X': X.astype(np.float32, copy=False)})[0]\n",
    "preds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "cbf704d7",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1.937151e-07"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.abs(preds - pp[:, [0, 2]]).max()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "735bb0d8",
   "metadata": {},
   "source": [
    "### Built-in wrapper\n",
    "\n",
    "As an alternative you can use wrapper that hide all the manipulations with `ONNX` and let you just call fit and predict. You can build wrapper from model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "5e5b46e7",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 100/100 [00:00<00:00, 1909.37it/s]\n"
     ]
    }
   ],
   "source": [
    "onnx_predictor = ONNXPredictor(\n",
    "    model, '../data/pb_model.onnx', \n",
    "    fltr=[0, 2], \n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "c506b652",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 4.71 s, sys: 156 ms, total: 4.86 s\n",
      "Wall time: 328 ms\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "array([[0.6264308 , 0.5388822 ],\n",
       "       [0.59586126, 0.56585   ],\n",
       "       [0.72726965, 0.49255225],\n",
       "       ...,\n",
       "       [0.5112887 , 0.4761739 ],\n",
       "       [0.67362005, 0.62851644],\n",
       "       [0.56341565, 0.41594115]], dtype=float32)"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%%time\n",
    "preds = onnx_predictor.predict(X)\n",
    "preds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "13f372d7",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1.7881393e-07"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.abs(preds - pp[:, [0, 2]]).max()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2e6ddf9f",
   "metadata": {},
   "source": [
    "***Note*** : You can not save `ONNXPredictor` object, since `onnxruntime.InferenceSession` is not pickable. Instead, to use it in the other session, you can restore it from `ONNX` model file. But note that in this case you will loose the information about postprocessing function, if it was not provided as `post_transform` to `ONNXPredictor`.\n",
    "\n",
    "First option, provide the post_transform to `ONNXPredictor`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "3bb457f2",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 100/100 [00:00<00:00, 2116.98it/s]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "1.937151e-07"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# build the predictor and save parsed as ../data/pb_model.onnx\n",
    "onnx_predictor = ONNXPredictor(\n",
    "    model, '../data/pb_model.onnx', \n",
    "    fltr=[0, 2], \n",
    "    post_transform='LOGISTIC' # provide the ONNX post_transform manually\n",
    ")\n",
    "\n",
    "# create new instance from ../data/pb_model.onnx\n",
    "onnx_predictor = ONNXPredictor.from_onnx('../data/pb_model.onnx')\n",
    "preds = onnx_predictor.predict(X)\n",
    "np.abs(preds - pp[:, [0, 2]]).max()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f1aef69c",
   "metadata": {},
   "source": [
    "Second, is to provide the python postprocessing function in the new session:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "88d0e7eb",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 100/100 [00:00<00:00, 2232.89it/s]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "1.7881393e-07"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# build the predictor and save parsed as ../data/pb_model.onnx\n",
    "onnx_predictor = ONNXPredictor(\n",
    "    model, '../data/pb_model.onnx', \n",
    "    fltr=[0, 2], \n",
    ")\n",
    "\n",
    "# create new instance from ../data/pb_model.onnx\n",
    "onnx_predictor = ONNXPredictor.from_onnx(\n",
    "    '../data/pb_model.onnx', \n",
    "    postprocess_fn=lambda x: 1 / (1 + np.exp(-x)) # provide py-boost postprocess_fn manually\n",
    ")\n",
    "preds = onnx_predictor.predict(X)\n",
    "np.abs(preds - pp[:, [0, 2]]).max()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "rapids-env",
   "language": "python",
   "name": "rapids-env"
  },
  "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.10.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}