tftrt_dense_switch.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# TFTRT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys,os\n",
    "sys.path.append(os.getcwd())\n",
    "while os.getcwd().split('/')[-1] != 'ml_sandbox': os.chdir('..')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Check Environment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Wed Feb 12 10:01:57 2020       \n",
      "+-----------------------------------------------------------------------------+\n",
      "| NVIDIA-SMI 440.33.01    Driver Version: 440.33.01    CUDA Version: 10.2     |\n",
      "|-------------------------------+----------------------+----------------------+\n",
      "| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |\n",
      "| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |\n",
      "|===============================+======================+======================|\n",
      "|   0  Tesla T4            Off  | 00000000:00:1E.0 Off |                    0 |\n",
      "| N/A   34C    P0    28W /  70W |   2083MiB / 15109MiB |      0%      Default |\n",
      "+-------------------------------+----------------------+----------------------+\n",
      "                                                                               \n",
      "+-----------------------------------------------------------------------------+\n",
      "| Processes:                                                       GPU Memory |\n",
      "|  GPU       PID   Type   Process name                             Usage      |\n",
      "|=============================================================================|\n",
      "|    0      2542      C   ...envs/ml_sandbox-YhSVM9Gx/bin/python3.6m   739MiB |\n",
      "|    0      2972      C   ...envs/ml_sandbox-YhSVM9Gx/bin/python3.6m  1333MiB |\n",
      "+-----------------------------------------------------------------------------+\n"
     ]
    }
   ],
   "source": [
    "!nvidia-smi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Tensorflow version:  2.1.0\n"
     ]
    }
   ],
   "source": [
    "import tensorflow as tf\n",
    "print(\"Tensorflow version: \", tf.version.VERSION)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Tensor Core GPU Present: True\n"
     ]
    }
   ],
   "source": [
    "# confirm gpu environment\n",
    "from tensorflow.python.client import device_lib\n",
    "\n",
    "def check_tensor_core_gpu_present():\n",
    "    local_device_protos = device_lib.list_local_devices()\n",
    "    for line in local_device_protos:\n",
    "        if \"compute capability\" in str(line):\n",
    "            compute_capability = float(line.physical_device_desc.split(\"compute capability: \")[-1])\n",
    "            if compute_capability>=7.0:\n",
    "                return True\n",
    "    \n",
    "print(\"Tensor Core GPU Present:\", check_tensor_core_gpu_present())\n",
    "tensor_core_gpu = check_tensor_core_gpu_present()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "compute_dtype = tf.float16"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Using TFTRT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "fp16 mode\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import time\n",
    "import numpy as np\n",
    "import tensorflow as tf\n",
    "from tensorflow.python.compiler.tensorrt import trt_convert as trt\n",
    "from tensorflow.python.saved_model import tag_constants\n",
    "if compute_dtype == tf.float16:\n",
    "    tf.keras.mixed_precision.experimental.set_policy(\"mixed_float16\")\n",
    "    print('fp16 mode')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Functional API"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MyModel(tf.keras.models.Model):\n",
    "    def __init__(self):\n",
    "        super(MyModel, self).__init__()\n",
    "        self.dense_layers = []\n",
    "        self.dense_layers.append(tf.keras.layers.Dense(128, activation='relu', input_shape=(1,)))\n",
    "        for i in range(1):\n",
    "            self.dense_layers.append(tf.keras.layers.Dense(128, activation='relu'))\n",
    "        self.dense_layers.append(tf.keras.layers.Dense(2, name='mypredict'))\n",
    "\n",
    "    @tf.function(input_signature=[{\n",
    "        'x1': tf.TensorSpec(shape=[None, 1], dtype=compute_dtype, name='x1'),\n",
    "        'x2': tf.TensorSpec(shape=[None, 1], dtype=compute_dtype, name='x2'),\n",
    "        'c': tf.TensorSpec(shape=[], dtype=tf.bool, name='c'),\n",
    "    }])\n",
    "    def call(self, inputs):\n",
    "        x1 = inputs['x1']\n",
    "        x2 = inputs['x2']\n",
    "        if inputs['c']:\n",
    "            for l in self.dense_layers:\n",
    "                x1 = l(x1)\n",
    "            return x1\n",
    "        else:\n",
    "            for l in self.dense_layers:\n",
    "                x2 = l(x2)\n",
    "            return x2\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = MyModel()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(2, 2), dtype=float16, numpy=\n",
       "array([[-0.1329,  0.0899],\n",
       "       [-0.2659,  0.1798]], dtype=float16)>"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#model.predict({\n",
    "model({\n",
    "    'x1': tf.constant([[1], [2]], compute_dtype),\n",
    "    'x2': tf.constant([[1], [2]], compute_dtype),\n",
    "    'c': True,\n",
    "})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_dir = 'tmp/models/dense/dense_switch'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:Skipping full serialization of Keras model <__main__.MyModel object at 0x7f09a030c048>, because its inputs are not defined.\n",
      "WARNING:tensorflow:From /home/ubuntu/.local/share/virtualenvs/ml_sandbox-YhSVM9Gx/lib/python3.6/site-packages/tensorflow_core/python/ops/resource_variable_ops.py:1786: calling BaseResourceVariable.__init__ (from tensorflow.python.ops.resource_variable_ops) with constraint is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "If using Keras pass *_constraint arguments to layers.\n",
      "INFO:tensorflow:Assets written to: tmp/models/dense/dense_switch/assets\n"
     ]
    }
   ],
   "source": [
    "tf.saved_model.save(model, model_dir)\n",
    "#model.save(model_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2020-02-12 10:02:10.839812: I tensorflow/stream_executor/platform/default/dso_loader.cc:44] Successfully opened dynamic library libnvinfer.so.6\n",
      "2020-02-12 10:02:10.841220: I tensorflow/stream_executor/platform/default/dso_loader.cc:44] Successfully opened dynamic library libnvinfer_plugin.so.6\n",
      "\n",
      "MetaGraphDef with tag-set: 'serve' contains the following SignatureDefs:\n",
      "\n",
      "signature_def['__saved_model_init_op']:\n",
      "  The given SavedModel SignatureDef contains the following input(s):\n",
      "  The given SavedModel SignatureDef contains the following output(s):\n",
      "    outputs['__saved_model_init_op'] tensor_info:\n",
      "        dtype: DT_INVALID\n",
      "        shape: unknown_rank\n",
      "        name: NoOp\n",
      "  Method name is: \n",
      "\n",
      "signature_def['serving_default']:\n",
      "  The given SavedModel SignatureDef contains the following input(s):\n",
      "    inputs['c'] tensor_info:\n",
      "        dtype: DT_BOOL\n",
      "        shape: ()\n",
      "        name: serving_default_c:0\n",
      "    inputs['x1'] tensor_info:\n",
      "        dtype: DT_HALF\n",
      "        shape: (-1, 1)\n",
      "        name: serving_default_x1:0\n",
      "    inputs['x2'] tensor_info:\n",
      "        dtype: DT_HALF\n",
      "        shape: (-1, 1)\n",
      "        name: serving_default_x2:0\n",
      "  The given SavedModel SignatureDef contains the following output(s):\n",
      "    outputs['output_0'] tensor_info:\n",
      "        dtype: DT_HALF\n",
      "        shape: (-1, 2)\n",
      "        name: StatefulPartitionedCall:0\n",
      "  Method name is: tensorflow/serving/predict\n",
      "WARNING:tensorflow:From /home/ubuntu/.local/share/virtualenvs/ml_sandbox-YhSVM9Gx/lib/python3.6/site-packages/tensorflow_core/python/ops/resource_variable_ops.py:1786: calling BaseResourceVariable.__init__ (from tensorflow.python.ops.resource_variable_ops) with constraint is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "If using Keras pass *_constraint arguments to layers.\n",
      "\n",
      "Defined Functions:\n",
      "  Function Name: 'call'\n",
      "    Option #1\n",
      "      Callable with:\n",
      "        Argument #1\n",
      "          DType: dict\n",
      "          Value: {'x1': TensorSpec(shape=(None, 1), dtype=tf.float16, name='x1'), 'x2': TensorSpec(shape=(None, 1), dtype=tf.float16, name='x2'), 'c': TensorSpec(shape=(), dtype=tf.bool, name='c')}\n"
     ]
    }
   ],
   "source": [
    "!saved_model_cli show --all --dir tmp/models/dense/dense_switch"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# raw TF"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<tf.Tensor: shape=(2, 2), dtype=float16, numpy=\n",
       "array([[-0.1329,  0.0899],\n",
       "       [-0.2659,  0.1798]], dtype=float16)>"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tf_model = tf.saved_model.load(model_dir)\n",
    "tf_model.call({\n",
    "    'x1': tf.constant([[1], [2]], compute_dtype),\n",
    "    'x2': tf.constant([[1], [2]], compute_dtype),\n",
    "    'c': True\n",
    "})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 1.16 s, sys: 212 ms, total: 1.37 s\n",
      "Wall time: 620 ms\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "for i in range(1000):\n",
    "    tf_model.call({\n",
    "        'x1': tf.constant([[1], [2]], compute_dtype),\n",
    "        'x2': tf.constant([[1], [2]], compute_dtype),\n",
    "        'c': True\n",
    "    })"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# TFTRT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Converting to TF-TRT...\n",
      "INFO:tensorflow:Linked TensorRT version: (6, 0, 1)\n",
      "INFO:tensorflow:Loaded TensorRT version: (6, 0, 1)\n",
      "==================\n",
      "input_name: StatefulPartitionedCall/cond/Func/StatefulPartitionedCall/input/_4/_17\n",
      "op: Switch\n",
      "node: name: \"StatefulPartitionedCall/cond/Func/StatefulPartitionedCall/input/_4/_17\"\n",
      "op: \"Switch\"\n",
      "input: \"Func/StatefulPartitionedCall/input/_4\"\n",
      "input: \"Func/StatefulPartitionedCall/input/_1\"\n",
      "attr {\n",
      "  key: \"T\"\n",
      "  value {\n",
      "    type: DT_RESOURCE\n",
      "  }\n",
      "}\n",
      "attr {\n",
      "  key: \"_class\"\n",
      "  value {\n",
      "    list {\n",
      "      s: \"Func/StatefulPartitionedCall/input/_4\"\n",
      "    }\n",
      "  }\n",
      "}\n",
      "experimental_debug_info {\n",
      "  original_node_names: \"Func/StatefulPartitionedCall/input/_4\"\n",
      "}\n",
      "\n"
     ]
    },
    {
     "ename": "ValueError",
     "evalue": "Cannot find the Placeholder op that is an input to the ReadVariableOp.",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mValueError\u001b[0m                                Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-15-73b06b66a015>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m      8\u001b[0m converter = trt.TrtGraphConverterV2(input_saved_model_dir=model_dir,\n\u001b[1;32m      9\u001b[0m                                     conversion_params=conversion_params)\n\u001b[0;32m---> 10\u001b[0;31m \u001b[0mconverter\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mconvert\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     11\u001b[0m \u001b[0mconverter\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msave\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0moutput_saved_model_dir\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mtrt_model_dir\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     12\u001b[0m \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m'Done Converting to TF-TRT'\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m~/.local/share/virtualenvs/ml_sandbox-YhSVM9Gx/lib/python3.6/site-packages/tensorflow_core/python/compiler/tensorrt/trt_convert.py\u001b[0m in \u001b[0;36mconvert\u001b[0;34m(self, calibration_input_fn)\u001b[0m\n\u001b[1;32m    978\u001b[0m                                   self._input_saved_model_tags)\n\u001b[1;32m    979\u001b[0m     \u001b[0mfunc\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_saved_model\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0msignatures\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0m_input_saved_model_signature_key\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 980\u001b[0;31m     \u001b[0mfrozen_func\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mconvert_to_constants\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mconvert_variables_to_constants_v2\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mfunc\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    981\u001b[0m     grappler_meta_graph_def = saver.export_meta_graph(\n\u001b[1;32m    982\u001b[0m         graph_def=frozen_func.graph.as_graph_def(), graph=frozen_func.graph)\n",
      "\u001b[0;32m~/.local/share/virtualenvs/ml_sandbox-YhSVM9Gx/lib/python3.6/site-packages/tensorflow_core/python/framework/convert_to_constants.py\u001b[0m in \u001b[0;36mconvert_variables_to_constants_v2\u001b[0;34m(func, lower_control_flow)\u001b[0m\n\u001b[1;32m    508\u001b[0m         \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m'op:'\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mname_to_node\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0minput_name\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mop\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    509\u001b[0m         \u001b[0mprint\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m'node:'\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mname_to_node\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0minput_name\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 510\u001b[0;31m         raise ValueError(\"Cannot find the Placeholder op that is an input \"\n\u001b[0m\u001b[1;32m    511\u001b[0m                          \"to the ReadVariableOp.\")\n\u001b[1;32m    512\u001b[0m       \u001b[0m_save_placeholder\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0minput_name\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnode\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mattr\u001b[0m\u001b[0;34m[\u001b[0m\u001b[0;34m\"dtype\"\u001b[0m\u001b[0;34m]\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;31mValueError\u001b[0m: Cannot find the Placeholder op that is an input to the ReadVariableOp."
     ]
    }
   ],
   "source": [
    "print('Converting to TF-TRT...')\n",
    "trt_model_dir = 'tmp/models/dense/dense_switch_trt'\n",
    "conversion_params = trt.DEFAULT_TRT_CONVERSION_PARAMS._replace(\n",
    "    precision_mode=trt.TrtPrecisionMode.FP16,\n",
    "    max_workspace_size_bytes=8000000000\n",
    ")\n",
    "\n",
    "converter = trt.TrtGraphConverterV2(input_saved_model_dir=model_dir,\n",
    "                                    conversion_params=conversion_params)\n",
    "converter.convert()\n",
    "converter.save(output_saved_model_dir=trt_model_dir)\n",
    "print('Done Converting to TF-TRT')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!saved_model_cli show --all --dir tmp/models/dense/dense_trtfp32"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "saved_model_loaded = tf.saved_model.load(trt_model_dir, tags=[tag_constants.SERVING])\n",
    "signature_keys = list(saved_model_loaded.signatures.keys())\n",
    "print(signature_keys)\n",
    "\n",
    "infer = saved_model_loaded.signatures['serving_default']\n",
    "print(infer.structured_outputs)\n",
    "\n",
    "y = infer(\n",
    "    x1=tf.constant([[1], [2]], compute_dtype),\n",
    "    x2=tf.constant([[1], [2]], compute_dtype),\n",
    "    c=tf.constant(True)\n",
    ")\n",
    "preds = y['output_0'].numpy()\n",
    "print(preds)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "for i in range(1000):\n",
    "    infer(\n",
    "        x1=tf.constant([[1], [2]], compute_dtype),\n",
    "        x2=tf.constant([[1], [2]], compute_dtype),\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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
}

I also meet the same error (tf2.2) . dou you kown how to solve the problem ?