onnx_quat.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import onnx\n",
    "from onnxruntime.quantization import quantize_dynamic, QuantType\n",
    "\n",
    "model_fp32 = 'assets/landmark_models/lm_model0_op11.onnx'\n",
    "model_quant = 'assets/landmark_models/lm_model0_opt_quat.onnx'\n",
    "model_quant = \"C:\\\\Users\\\\xuhao\\\\Develop\\\\build-FlightAgentX-Desktop_Qt_5_14_2_MSVC2017_32bit-Release\\\\release\\\\assets\\\\landmark_models\\\\lm_model0_opt_quat.onnx\"\n",
    "calibration_dataset_path = \"D:\\\\data\\\\LS3D-W.tar\\\\LS3D-W\\\\Menpo-3D\"\n",
    "#quantized_model = quantize_dynamic(model_fp32, model_quant, weight_type=QuantType.QUInt8)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Collect 17910 files\n",
      "input 8955\n"
     ]
    },
    {
     "ename": "FileNotFoundError",
     "evalue": "[Errno 2] No such file or directory: 'assets/landmark_models/lm_model0_op11.onnx'",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mFileNotFoundError\u001b[0m                         Traceback (most recent call last)",
      "\u001b[1;32m<ipython-input-2-8d107f498f68>\u001b[0m in \u001b[0;36m<module>\u001b[1;34m()\u001b[0m\n\u001b[0;32m     82\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m     83\u001b[0m \u001b[0mdr\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mResNet50DataReader\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mcalibration_dataset_path\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mmodel_fp32\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m---> 84\u001b[1;33m \u001b[0mquantize_static\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mmodel_fp32\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mmodel_quant\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mdr\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mweight_type\u001b[0m\u001b[1;33m=\u001b[0m\u001b[0mQuantType\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mQUInt8\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m     85\u001b[0m \u001b[0mprint\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;34m'Calibrated and quantized model saved.'\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32m~\\Anaconda3\\lib\\site-packages\\onnxruntime\\quantization\\quantize.py\u001b[0m in \u001b[0;36mquantize_static\u001b[1;34m(model_input, model_output, calibration_data_reader, op_types_to_quantize, per_channel, reduce_range, activation_type, weight_type, nodes_to_quantize, nodes_to_exclude, use_external_data_format)\u001b[0m\n\u001b[0;32m    169\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    170\u001b[0m     quantization_params_dict = calibrate(model_input, calibration_data_reader, op_types_to_quantize, nodes_to_quantize,\n\u001b[1;32m--> 171\u001b[1;33m                                          nodes_to_exclude)\n\u001b[0m\u001b[0;32m    172\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    173\u001b[0m     quantizer = ONNXQuantizer(\n",
      "\u001b[1;32m~\\Anaconda3\\lib\\site-packages\\onnxruntime\\quantization\\calibrate.py\u001b[0m in \u001b[0;36mcalibrate\u001b[1;34m(model_path, data_reader, op_types, black_nodes, white_nodes, augmented_model_path)\u001b[0m\n\u001b[0;32m    259\u001b[0m     \u001b[0mdict_for_quantization\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mcalibrater\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mget_intermediate_outputs\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    260\u001b[0m     \u001b[1;31m#4. generate quantization parameters dict\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m--> 261\u001b[1;33m     \u001b[0mquantization_params_dict\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mcalibrater\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mcalculate_quantization_params\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mdict_for_quantization\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m    262\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    263\u001b[0m     \u001b[0mprint\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;34m\"Calibrated,quantized parameters calculated and returned.\"\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32m~\\Anaconda3\\lib\\site-packages\\onnxruntime\\quantization\\calibrate.py\u001b[0m in \u001b[0;36mcalculate_quantization_params\u001b[1;34m(self, quantization_thresholds)\u001b[0m\n\u001b[0;32m    218\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    219\u001b[0m         \u001b[0mquantization_params\u001b[0m \u001b[1;33m=\u001b[0m \u001b[1;33m{\u001b[0m\u001b[1;33m}\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m--> 220\u001b[1;33m         \u001b[0mmodel\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0monnx\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mload\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mmodel_path\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m    221\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    222\u001b[0m         \u001b[0mself\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0m_get_input_name_to_nodes\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mmodel\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32m~\\Anaconda3\\lib\\site-packages\\onnx\\__init__.py\u001b[0m in \u001b[0;36mload_model\u001b[1;34m(f, format, load_external_data)\u001b[0m\n\u001b[0;32m    112\u001b[0m     \u001b[0mLoaded\u001b[0m \u001b[1;32min\u001b[0m\u001b[1;33m-\u001b[0m\u001b[0mmemory\u001b[0m \u001b[0mModelProto\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    113\u001b[0m     '''\n\u001b[1;32m--> 114\u001b[1;33m     \u001b[0ms\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0m_load_bytes\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mf\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m    115\u001b[0m     \u001b[0mmodel\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mload_model_from_string\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0ms\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mformat\u001b[0m\u001b[1;33m=\u001b[0m\u001b[0mformat\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    116\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32m~\\Anaconda3\\lib\\site-packages\\onnx\\__init__.py\u001b[0m in \u001b[0;36m_load_bytes\u001b[1;34m(f)\u001b[0m\n\u001b[0;32m     28\u001b[0m         \u001b[0ms\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mcast\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mIO\u001b[0m\u001b[1;33m[\u001b[0m\u001b[0mbytes\u001b[0m\u001b[1;33m]\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mf\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mread\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m     29\u001b[0m     \u001b[1;32melse\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m---> 30\u001b[1;33m         \u001b[1;32mwith\u001b[0m \u001b[0mopen\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mcast\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mText\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mf\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m,\u001b[0m \u001b[1;34m'rb'\u001b[0m\u001b[1;33m)\u001b[0m \u001b[1;32mas\u001b[0m \u001b[0mreadable\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m     31\u001b[0m             \u001b[0ms\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mreadable\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mread\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m     32\u001b[0m     \u001b[1;32mreturn\u001b[0m \u001b[0ms\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;31mFileNotFoundError\u001b[0m: [Errno 2] No such file or directory: 'assets/landmark_models/lm_model0_op11.onnx'"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import sys\n",
    "import numpy as np\n",
    "import re\n",
    "import abc\n",
    "import subprocess\n",
    "import json\n",
    "import cv2\n",
    "\n",
    "import onnx\n",
    "import onnxruntime\n",
    "from onnx import helper, TensorProto, numpy_helper\n",
    "from onnxruntime.quantization import quantize_static, calibrate, CalibrationDataReader\n",
    "\n",
    "##Quat for lm model\n",
    "\n",
    "class ResNet50DataReader(CalibrationDataReader):\n",
    "    def __init__(self, calibration_image_folder, augmented_model_path='augmented_model.onnx', size_limit=0):\n",
    "        self.image_folder = calibration_image_folder\n",
    "        self.augmented_model_path = augmented_model_path\n",
    "        self.preprocess_flag = True\n",
    "        self.enum_data_dicts = []\n",
    "        self.datasize = 0\n",
    "        \n",
    "\n",
    "        self.size_limit = size_limit\n",
    "\n",
    "\n",
    "    def get_next(self):\n",
    "        if self.preprocess_flag:\n",
    "            self.preprocess_flag = False\n",
    "            session = onnxruntime.InferenceSession(self.augmented_model_path, None)\n",
    "            (_, height, width, _) = session.get_inputs()[0].shape\n",
    "            nhwc_data_list = preprocess_func(self.image_folder, 224, 224, size_limit=self.size_limit)\n",
    "            input_name = session.get_inputs()[0].name\n",
    "            print(input_name, len(nhwc_data_list))\n",
    "            self.datasize = len(nhwc_data_list)\n",
    "            self.enum_data_dicts = iter([{input_name: nhwc_data} for nhwc_data in nhwc_data_list])\n",
    "        return next(self.enum_data_dicts, None)\n",
    "\n",
    "\n",
    "def preprocess_func(images_folder, height, width, size_limit=0):\n",
    "    '''\n",
    "    Loads a batch of images and preprocess them\n",
    "    parameter images_folder: path to folder storing images\n",
    "    parameter height: image height in pixels\n",
    "    parameter width: image width in pixels\n",
    "    parameter size_limit: number of images to load. Default is 0 which means all images are picked.\n",
    "    return: list of matrices characterizing multiple images\n",
    "    '''\n",
    "    mean = np.float32(np.array([0.485, 0.456, 0.406]))\n",
    "    std = np.float32(np.array([0.229, 0.224, 0.225]))\n",
    "    mean = mean / std\n",
    "    std = std * 255.0\n",
    "    mean = - mean\n",
    "    std = 1.0 / std\n",
    "    mean_32 = np.tile(mean, [32, 32, 1])\n",
    "    std_32 = np.tile(std, [32, 32, 1])\n",
    "    mean_224 = np.tile(mean, [224, 224, 1])\n",
    "    std_224 = np.tile(std, [224, 224, 1])\n",
    "        \n",
    "    image_names = os.listdir(images_folder)\n",
    "    if size_limit > 0 and len(image_names) >= size_limit:\n",
    "        batch_filenames = [image_names[i] for i in range(size_limit)]\n",
    "    else:\n",
    "        batch_filenames = image_names\n",
    "    unconcatenated_batch_data = []\n",
    "    print(f\"Collect {len(batch_filenames)} files\")\n",
    "    for image_name in batch_filenames:\n",
    "        if image_name.endswith(\".jpg\"):\n",
    "            image_filepath = images_folder + '/' + image_name\n",
    "            frame = cv2.imread(image_filepath, cv2.IMREAD_COLOR)\n",
    "            frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)\n",
    "            im = cv2.resize(frame, (width, height), interpolation=cv2.INTER_LINEAR)[:,:,::-1]\n",
    "            im = im * std_224 + mean_224\n",
    "            im = np.expand_dims(im, 0)\n",
    "            im = np.transpose(im, (0,3,1,2))\n",
    "            unconcatenated_batch_data.append(im)\n",
    "    batch_data = np.concatenate(np.expand_dims(unconcatenated_batch_data, axis=0), axis=0)\n",
    "    return batch_data\n",
    "\n",
    "\n",
    "dr = ResNet50DataReader(calibration_dataset_path, model_fp32)\n",
    "quantize_static(model_fp32, model_quant, dr, weight_type=QuantType.QUInt8)\n",
    "print('Calibrated and quantized model saved.')\n",
    "\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.6.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}