onnx-to-triton.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "SPDX-FileCopyrightText: Copyright (c) 2021 NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n",
    "\n",
    "SPDX-License-Identifier: MIT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import onnx\n",
    "from onnx import TensorProto\n",
    "from onnx.helper import tensor_dtype_to_np_dtype\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "def type_conv(t): # see https://github.com/onnx/onnx/blob/master/onnx/mapping.py#L12\n",
    "    dt = tensor_dtype_to_np_dtype(t)\n",
    "    dt = dt.name.upper()\n",
    "    dt = dt.replace(\"FLOAT\", \"FP\")\n",
    "    return f\"TYPE_{dt}\"\n",
    "\n",
    "def node_dims(node, include_batch):\n",
    "    shape = node.type.tensor_type.shape.dim\n",
    "    dims = []\n",
    "    for i, d in enumerate(shape):\n",
    "        if not include_batch and i == 0: continue\n",
    "        if d.HasField(\"dim_param\"): dims += [-1]\n",
    "        elif d.HasField(\"dim_value\"): dims += [d.dim_value]\n",
    "        else: raise(Exception(f\"Unknown dimension {d}\"))\n",
    "    s = \", \".join(str(d) for d in dims)\n",
    "    if not include_batch and not len(dims):\n",
    "        return (\"dims: [ 1 ]\\n\" + \n",
    "               \"      reshape: { shape: [ ] }\")\n",
    "    s = f\"dims: [ {s} ]\"\n",
    "    return s\n",
    "\n",
    "def node_str(node, include_batch=False):\n",
    "    return \\\n",
    "f\"\"\"\n",
    "    {{\n",
    "      name: \"{node.name}\"\n",
    "      data_type: {type_conv(node.type.tensor_type.elem_type)}\n",
    "      {node_dims(node, include_batch)}\n",
    "    }}\"\"\"\n",
    "\n",
    "def graph_str(graph, include_batch=False):\n",
    "    s = \"\"\n",
    "    s += \"input [\"\n",
    "    t = [node_str(n, include_batch) for n in graph.input]\n",
    "    s += \",\\n\".join(t)\n",
    "    s += \"\\n]\\n\\n\"\n",
    "\n",
    "    s += \"output [\"\n",
    "    t = [node_str(n, include_batch) for n in graph.output]\n",
    "    s += \",\\n\".join(t)\n",
    "    s += \"\\n]\"\n",
    "\n",
    "    return s\n",
    "\n",
    "def header_str(model_onnx, max_batch=None):\n",
    "    s = 'platform: \"onnxruntime_onnx\"\\n'\n",
    "    if max_batch is not None: s+= f'max_batch_size: {max_batch}\\n'\n",
    "    s += \"\\n\"\n",
    "    return s\n",
    "\n",
    "def onnx_to_str(f, max_batch=None):\n",
    "    model_onnx = onnx.load(f)\n",
    "    graph = model_onnx.graph\n",
    "    return header_str(model_onnx, max_batch) + graph_str(graph, include_batch=max_batch is None)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "platform: \"onnxruntime_onnx\"\n",
      "max_batch_size: 16\n",
      "\n",
      "input [\n",
      "    {\n",
      "      name: \"text\"\n",
      "      data_type: TYPE_INT64\n",
      "      dims: [ -1 ]\n",
      "    }\n",
      "]\n",
      "\n",
      "output [\n",
      "    {\n",
      "      name: \"spect\"\n",
      "      data_type: TYPE_FP32\n",
      "      dims: [ 80, -1 ]\n",
      "    },\n",
      "\n",
      "    {\n",
      "      name: \"num_frames\"\n",
      "      data_type: TYPE_INT64\n",
      "      dims: [ 1 ]\n",
      "      reshape: { shape: [ ] }\n",
      "    },\n",
      "\n",
      "    {\n",
      "      name: \"durs_predicted\"\n",
      "      data_type: TYPE_FP32\n",
      "      dims: [ -1 ]\n",
      "    },\n",
      "\n",
      "    {\n",
      "      name: \"log_durs_predicted\"\n",
      "      data_type: TYPE_FP32\n",
      "      dims: [ -1 ]\n",
      "    },\n",
      "\n",
      "    {\n",
      "      name: \"pitch_predicted\"\n",
      "      data_type: TYPE_FP32\n",
      "      dims: [ -1 ]\n",
      "    }\n",
      "]\n"
     ]
    }
   ],
   "source": [
    "print(onnx_to_str(\"fastpitch.onnx\", 16))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(\"model_repository/fastpitch/config.pbtxt\", \"w\") as cfg:\n",
    "    cfg.write(onnx_to_str(\"fastpitch.onnx\", 128))\n",
    "\n",
    "with open(\"model_repository/hifigan/config.pbtxt\", \"w\") as cfg:\n",
    "    cfg.write(onnx_to_str(\"hifigan.onnx\", 128))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "dim_param: \"text_dynamic_axes_1\""
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "node.type.tensor_type.shape.dim[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "!cp onnx-to-triton.ipynb /ext_home/workspace_ru/code"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "98b0a9b7b4eaaa670588a142fd0a9b87eaafe866f1db4228be72b4211d12040f"
  },
  "kernelspec": {
   "display_name": "Python 3.8.10 64-bit ('base': conda)",
   "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.8.10"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}