├── requirements.txt ├── data ├── img.png ├── dogs.jpg ├── elephant.jpg └── man-dog.jpg ├── .gitignore ├── LICENSE ├── README.md ├── Step2_keras-jetson-ImageNet-predict.ipynb ├── Step1_Colab_TensorRT.ipynb └── Step1_Object_detection_Colab_TensorRT.ipynb /requirements.txt: -------------------------------------------------------------------------------- 1 | tensorflow -------------------------------------------------------------------------------- /data/img.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Tony607/tf_jetson_nano/HEAD/data/img.png -------------------------------------------------------------------------------- /data/dogs.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Tony607/tf_jetson_nano/HEAD/data/dogs.jpg -------------------------------------------------------------------------------- /data/elephant.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Tony607/tf_jetson_nano/HEAD/data/elephant.jpg -------------------------------------------------------------------------------- /data/man-dog.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Tony607/tf_jetson_nano/HEAD/data/man-dog.jpg -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | *.pyc 2 | *.hdf5 3 | *.ipynb_checkpoints 4 | *.p 5 | *.h5 6 | *.HDF5 7 | __pycache__ 8 | # Ignore model directory. 9 | model/ 10 | .vscode/ -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | COPYRIGHT 2 | 3 | All contributions by François Chollet: 4 | Copyright (c) 2015 - 2018, François Chollet. 5 | All rights reserved. 6 | 7 | All contributions by Google: 8 | Copyright (c) 2015 - 2018, Google, Inc. 9 | All rights reserved. 10 | 11 | All contributions by Microsoft: 12 | Copyright (c) 2017 - 2018, Microsoft, Inc. 13 | All rights reserved. 14 | 15 | All other contributions: 16 | Copyright (c) 2015 - 2018, the respective contributors. 17 | All rights reserved. 18 | 19 | Each contributor holds copyright over their respective contributions. 20 | The project versioning (Git) records all such contribution source information. 21 | 22 | LICENSE 23 | 24 | The MIT License (MIT) 25 | 26 | Permission is hereby granted, free of charge, to any person obtaining a copy 27 | of this software and associated documentation files (the "Software"), to deal 28 | in the Software without restriction, including without limitation the rights 29 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 30 | copies of the Software, and to permit persons to whom the Software is 31 | furnished to do so, subject to the following conditions: 32 | 33 | The above copyright notice and this permission notice shall be included in all 34 | copies or substantial portions of the Software. 35 | 36 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 37 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 38 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 39 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 40 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 41 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 42 | SOFTWARE. 43 | 44 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Run Keras/Tensorflow model on Jetson Nano 2 | 3 | 4 | ## How to Run 5 | Require [Python 3.5+](https://www.python.org/ftp/python/3.6.7/python-3.6.7.exe) and [Jupyter notebook](https://jupyter.readthedocs.io/en/latest/install.html) installed 6 | ### Clone or download this repo 7 | ``` 8 | git clone https://github.com/Tony607/tf_jetson_nano 9 | ``` 10 | 11 | ### Install required libraries for your development machine 12 | `pip3 install -r requirements.txt` 13 | ## A. The Keras Image classification model 14 | [How to run Keras model on Jetson Nano](https://www.dlology.com/blog/how-to-run-keras-model-on-jetson-nano/) | DLology Blog 15 | ### Step1: Convert Keras model into TensorRT model 16 | On development machine or [Google Colab](https://colab.research.google.com/github/Tony607/tf_jetson_nano/blob/master/Step1_Colab_TensorRT.ipynb) 17 | 18 | 19 | To run locally, start a terminal, then run, 20 | ``` 21 | jupyter notebook 22 | ``` 23 | 24 | In the opened browser window open 25 | ``` 26 | Step1_Colab_TensorRT.ipynb 27 | ``` 28 | 29 | ### Step2: Make prediction (On Jetson Nano) 30 | To run notebook on Jetson Nano and make it accessible from another machine, in its terminal run, 31 | ``` 32 | jupyter notebook --ip=0.0.0.0 33 | ``` 34 | In another machine's browser, open the notebook, 35 | ``` 36 | http://:8888/notebooks/Step2_keras-jetson-ImageNet-predict.ipynb 37 | ``` 38 | ## B.Tensorflow object detection model 39 | [How to run TensorFlow Object Detection model on Jetson Nano](https://www.dlology.com/blog/how-to-run-tensorflow-object-detection-model-on-jetson-nano/) | DLology Blog 40 | 41 | ### Step1: Convert Tensorflow object detection model into TensorRT model 42 | On development machine or [Google Colab Notebook](https://colab.research.google.com/github/Tony607/tf_jetson_nano/blob/master/Step1_Object_detection_Colab_TensorRT.ipynb) 43 | 44 | To run locally, start a terminal, then run, 45 | ``` 46 | jupyter notebook 47 | ``` 48 | 49 | In the opened browser window open 50 | ``` 51 | Step1_Object_detection_Colab_TensorRT.ipynb 52 | ``` 53 | 54 | ### Step2: Make prediction (On Jetson Nano) 55 | To run notebook on Jetson Nano and make it accessible from another machine, in its terminal run, 56 | ``` 57 | jupyter notebook --ip=0.0.0.0 58 | ``` 59 | In another machine's browser, open the notebook, 60 | ``` 61 | http://:8888/notebooks/Step2_jetson-object-detection-predict.ipynb 62 | ``` -------------------------------------------------------------------------------- /Step2_keras-jetson-ImageNet-predict.ipynb: -------------------------------------------------------------------------------- 1 | { 2 | "cells": [ 3 | { 4 | "cell_type": "markdown", 5 | "metadata": {}, 6 | "source": [ 7 | "# Step 2: Make prediction with TensorRT graph.\n", 8 | "\n", 9 | "Run this notebook on Jetson Nano.\n", 10 | "**For more detail, checkout [How to run Keras model on Jetson Nano](https://www.dlology.com/blog/how-to-run-keras-model-on-jetson-nano/) | DLology Blog**" 11 | ] 12 | }, 13 | { 14 | "cell_type": "markdown", 15 | "metadata": {}, 16 | "source": [ 17 | "Input and output nodes names printed in the output during freezing Keras model in previous step." 18 | ] 19 | }, 20 | { 21 | "cell_type": "code", 22 | "execution_count": 1, 23 | "metadata": {}, 24 | "outputs": [], 25 | "source": [ 26 | "output_names = ['Logits/Softmax']\n", 27 | "input_names = ['input_1']" 28 | ] 29 | }, 30 | { 31 | "cell_type": "markdown", 32 | "metadata": {}, 33 | "source": [ 34 | "## Load the `TensorRT` frozen graph `.pb` file\n", 35 | "This might take a while." 36 | ] 37 | }, 38 | { 39 | "cell_type": "code", 40 | "execution_count": 3, 41 | "metadata": {}, 42 | "outputs": [ 43 | { 44 | "name": "stdout", 45 | "output_type": "stream", 46 | "text": [ 47 | "WARNING:tensorflow:From :6: FastGFile.__init__ (from tensorflow.python.platform.gfile) is deprecated and will be removed in a future version.\n", 48 | "Instructions for updating:\n", 49 | "Use tf.gfile.GFile.\n" 50 | ] 51 | } 52 | ], 53 | "source": [ 54 | "import tensorflow as tf\n", 55 | "import tensorflow.contrib.tensorrt as trt\n", 56 | "import numpy as np\n", 57 | "\n", 58 | "def get_frozen_graph(graph_file):\n", 59 | " \"\"\"Read Frozen Graph file from disk.\"\"\"\n", 60 | " with tf.gfile.FastGFile(graph_file, \"rb\") as f:\n", 61 | " graph_def = tf.GraphDef()\n", 62 | " graph_def.ParseFromString(f.read())\n", 63 | " return graph_def\n", 64 | "\n", 65 | "\n", 66 | "trt_graph = get_frozen_graph('./model/trt_graph.pb')" 67 | ] 68 | }, 69 | { 70 | "cell_type": "markdown", 71 | "metadata": {}, 72 | "source": [ 73 | "## Create session and load graph" 74 | ] 75 | }, 76 | { 77 | "cell_type": "code", 78 | "execution_count": 4, 79 | "metadata": {}, 80 | "outputs": [], 81 | "source": [ 82 | "# Create session and load graph\n", 83 | "tf_config = tf.ConfigProto()\n", 84 | "tf_config.gpu_options.allow_growth = True\n", 85 | "tf_sess = tf.Session(config=tf_config)\n", 86 | "tf.import_graph_def(trt_graph, name='')" 87 | ] 88 | }, 89 | { 90 | "cell_type": "code", 91 | "execution_count": 43, 92 | "metadata": {}, 93 | "outputs": [ 94 | { 95 | "name": "stdout", 96 | "output_type": "stream", 97 | "text": [ 98 | "image_size: [224, 224, 3]\n" 99 | ] 100 | } 101 | ], 102 | "source": [ 103 | "# Get graph input size\n", 104 | "for node in trt_graph.node:\n", 105 | " if 'input_' in node.name:\n", 106 | " size = node.attr['shape'].shape\n", 107 | " image_size = [size.dim[i].size for i in range(1, 4)]\n", 108 | " break\n", 109 | "print(\"image_size: {}\".format(image_size))" 110 | ] 111 | }, 112 | { 113 | "cell_type": "code", 114 | "execution_count": 59, 115 | "metadata": {}, 116 | "outputs": [ 117 | { 118 | "name": "stdout", 119 | "output_type": "stream", 120 | "text": [ 121 | "input_tensor_name: input_1:0\n", 122 | "output_tensor_name: Logits/Softmax:0\n" 123 | ] 124 | } 125 | ], 126 | "source": [ 127 | "input_tensor_name = input_names[0] + \":0\"\n", 128 | "output_tensor_name = output_names[0] + \":0\"\n", 129 | "\n", 130 | "print(\"input_tensor_name: {}\\noutput_tensor_name: {}\".format(\n", 131 | " input_tensor_name, output_tensor_name))\n", 132 | "\n", 133 | "output_tensor = tf_sess.graph.get_tensor_by_name(output_tensor_name)" 134 | ] 135 | }, 136 | { 137 | "cell_type": "markdown", 138 | "metadata": {}, 139 | "source": [ 140 | "## Predict a real image" 141 | ] 142 | }, 143 | { 144 | "cell_type": "code", 145 | "execution_count": 60, 146 | "metadata": {}, 147 | "outputs": [], 148 | "source": [ 149 | "from tensorflow.keras.preprocessing import image\n", 150 | "from tensorflow.keras.applications.mobilenet_v2 import preprocess_input, decode_predictions" 151 | ] 152 | }, 153 | { 154 | "cell_type": "code", 155 | "execution_count": 61, 156 | "metadata": {}, 157 | "outputs": [], 158 | "source": [ 159 | "# Optional image to test model prediction.\n", 160 | "img_path = './data/elephant.jpg'\n", 161 | "\n", 162 | "img = image.load_img(img_path, target_size=image_size[:2])\n", 163 | "x = image.img_to_array(img)\n", 164 | "x = np.expand_dims(x, axis=0)\n", 165 | "x = preprocess_input(x)" 166 | ] 167 | }, 168 | { 169 | "cell_type": "code", 170 | "execution_count": 62, 171 | "metadata": {}, 172 | "outputs": [ 173 | { 174 | "data": { 175 | "text/plain": [ 176 | "(1, 1000)" 177 | ] 178 | }, 179 | "execution_count": 62, 180 | "metadata": {}, 181 | "output_type": "execute_result" 182 | } 183 | ], 184 | "source": [ 185 | "feed_dict = {\n", 186 | " input_tensor_name: x\n", 187 | "}\n", 188 | "preds = tf_sess.run(output_tensor, feed_dict)\n", 189 | "\n", 190 | "preds.shape" 191 | ] 192 | }, 193 | { 194 | "cell_type": "code", 195 | "execution_count": 63, 196 | "metadata": {}, 197 | "outputs": [ 198 | { 199 | "name": "stdout", 200 | "output_type": "stream", 201 | "text": [ 202 | "Predicted: [('n02504458', 'African_elephant', 0.6990039), ('n01871265', 'tusker', 0.16008943), ('n02504013', 'Indian_elephant', 0.03917188)]\n" 203 | ] 204 | } 205 | ], 206 | "source": [ 207 | "# decode the results into a list of tuples (class, description, probability)\n", 208 | "# (one such list for each sample in the batch)\n", 209 | "print('Predicted:', decode_predictions(preds, top=3)[0])" 210 | ] 211 | }, 212 | { 213 | "cell_type": "markdown", 214 | "metadata": {}, 215 | "source": [ 216 | "## Benchmark inference speed" 217 | ] 218 | }, 219 | { 220 | "cell_type": "code", 221 | "execution_count": 64, 222 | "metadata": {}, 223 | "outputs": [ 224 | { 225 | "name": "stdout", 226 | "output_type": "stream", 227 | "text": [ 228 | "average(sec):0.04,fps:27.18\n" 229 | ] 230 | } 231 | ], 232 | "source": [ 233 | "import time\n", 234 | "times = []\n", 235 | "for i in range(20):\n", 236 | " start_time = time.time()\n", 237 | " one_prediction = tf_sess.run(output_tensor, feed_dict)\n", 238 | " delta = (time.time() - start_time)\n", 239 | " times.append(delta)\n", 240 | "mean_delta = np.array(times).mean()\n", 241 | "fps = 1 / mean_delta\n", 242 | "print('average(sec):{:.2f},fps:{:.2f}'.format(mean_delta, fps))" 243 | ] 244 | }, 245 | { 246 | "cell_type": "markdown", 247 | "metadata": {}, 248 | "source": [ 249 | "## Close the session to release resources" 250 | ] 251 | }, 252 | { 253 | "cell_type": "code", 254 | "execution_count": null, 255 | "metadata": {}, 256 | "outputs": [], 257 | "source": [ 258 | "tf_sess.close()" 259 | ] 260 | } 261 | ], 262 | "metadata": { 263 | "kernelspec": { 264 | "display_name": "Python 3", 265 | "language": "python", 266 | "name": "python3" 267 | }, 268 | "language_info": { 269 | "codemirror_mode": { 270 | "name": "ipython", 271 | "version": 3 272 | }, 273 | "file_extension": ".py", 274 | "mimetype": "text/x-python", 275 | "name": "python", 276 | "nbconvert_exporter": "python", 277 | "pygments_lexer": "ipython3", 278 | "version": "3.5.2" 279 | }, 280 | "toc": { 281 | "base_numbering": 1, 282 | "nav_menu": {}, 283 | "number_sections": true, 284 | "sideBar": true, 285 | "skip_h1_title": false, 286 | "title_cell": "Table of Contents", 287 | "title_sidebar": "Contents", 288 | "toc_cell": false, 289 | "toc_position": {}, 290 | "toc_section_display": true, 291 | "toc_window_display": false 292 | } 293 | }, 294 | "nbformat": 4, 295 | "nbformat_minor": 2 296 | } 297 | -------------------------------------------------------------------------------- /Step1_Colab_TensorRT.ipynb: -------------------------------------------------------------------------------- 1 | { 2 | "nbformat": 4, 3 | "nbformat_minor": 0, 4 | "metadata": { 5 | "colab": { 6 | "name": "Step1_Colab_TensorRT.ipynb", 7 | "version": "0.3.2", 8 | "provenance": [], 9 | "collapsed_sections": [], 10 | "include_colab_link": true 11 | }, 12 | "kernelspec": { 13 | "name": "python3", 14 | "display_name": "Python 3" 15 | }, 16 | "accelerator": "GPU" 17 | }, 18 | "cells": [ 19 | { 20 | "cell_type": "markdown", 21 | "metadata": { 22 | "id": "view-in-github", 23 | "colab_type": "text" 24 | }, 25 | "source": [ 26 | "\"Open" 27 | ] 28 | }, 29 | { 30 | "metadata": { 31 | "id": "QQK7KlhRBfnf", 32 | "colab_type": "text" 33 | }, 34 | "cell_type": "markdown", 35 | "source": [ 36 | "# Step 1 : Convert Keras model into TensorRT model\n", 37 | "**For more detail, checkout [How to run Keras model on Jetson Nano](https://www.dlology.com/blog/how-to-run-keras-model-on-jetson-nano/) | DLology Blog**" 38 | ] 39 | }, 40 | { 41 | "metadata": { 42 | "id": "ITUc4J7etACS", 43 | "colab_type": "code", 44 | "colab": {} 45 | }, 46 | "cell_type": "code", 47 | "source": [ 48 | "import tensorflow.contrib.tensorrt as trt" 49 | ], 50 | "execution_count": 0, 51 | "outputs": [] 52 | }, 53 | { 54 | "metadata": { 55 | "id": "VsZgX4XeuNkt", 56 | "colab_type": "code", 57 | "colab": {} 58 | }, 59 | "cell_type": "code", 60 | "source": [ 61 | "# Download an image for prediction.\n", 62 | "\n", 63 | "!wget --quiet https://raw.githubusercontent.com/Tony607/tf_jetson_nano/master/data/elephant.jpg" 64 | ], 65 | "execution_count": 0, 66 | "outputs": [] 67 | }, 68 | { 69 | "metadata": { 70 | "id": "fWygvIyctpeI", 71 | "colab_type": "code", 72 | "outputId": "d2dea334-f704-4658-a495-c3c9d559e472", 73 | "colab": { 74 | "base_uri": "https://localhost:8080/", 75 | "height": 110 76 | } 77 | }, 78 | "cell_type": "code", 79 | "source": [ 80 | "\n", 81 | "import tensorflow as tf\n", 82 | "from tensorflow.keras.applications.mobilenet_v2 import MobileNetV2 as Net\n", 83 | "from tensorflow.keras.preprocessing import image\n", 84 | "from tensorflow.keras.applications.mobilenet_v2 import preprocess_input, decode_predictions\n", 85 | "import numpy as np\n", 86 | "import os\n", 87 | "\n", 88 | "# Optional image to test model prediction.\n", 89 | "img_path = './elephant.jpg'\n", 90 | "model_path = './model'\n", 91 | "\n", 92 | "os.makedirs(model_path, exist_ok=True)\n", 93 | "# Path to save the model h5 file.\n", 94 | "model_fname = os.path.join(model_path, 'model.h5')\n", 95 | "\n", 96 | "os.makedirs(model_path, exist_ok=True)\n", 97 | "\n", 98 | "img_height = 224\n", 99 | "\n", 100 | "model = Net(weights='imagenet')\n", 101 | "\n", 102 | "\n", 103 | "# Load the image for prediction.\n", 104 | "img = image.load_img(img_path, target_size=(img_height, img_height))\n", 105 | "x = image.img_to_array(img)\n", 106 | "x = np.expand_dims(x, axis=0)\n", 107 | "x = preprocess_input(x)\n", 108 | "\n", 109 | "preds = model.predict(x)\n", 110 | "# decode the results into a list of tuples (class, description, probability)\n", 111 | "# (one such list for each sample in the batch)\n", 112 | "print('Predicted:', decode_predictions(preds, top=3)[0])\n", 113 | "# Predicted: [(u'n02504013', u'Indian_elephant', 0.82658225), (u'n01871265', u'tusker', 0.1122357), (u'n02504458', u'African_elephant', 0.061040461)]\n", 114 | "\n", 115 | "# Save the h5 file to path specified.\n", 116 | "model.save(model_fname)" 117 | ], 118 | "execution_count": 3, 119 | "outputs": [ 120 | { 121 | "output_type": "stream", 122 | "text": [ 123 | "WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensorflow/python/ops/resource_variable_ops.py:435: colocate_with (from tensorflow.python.framework.ops) is deprecated and will be removed in a future version.\n", 124 | "Instructions for updating:\n", 125 | "Colocations handled automatically by placer.\n", 126 | "Predicted: [('n02504458', 'African_elephant', 0.705134), ('n01871265', 'tusker', 0.1551433), ('n02504013', 'Indian_elephant', 0.039011613)]\n" 127 | ], 128 | "name": "stdout" 129 | } 130 | ] 131 | }, 132 | { 133 | "metadata": { 134 | "id": "OY9k1o9iulZF", 135 | "colab_type": "text" 136 | }, 137 | "cell_type": "markdown", 138 | "source": [ 139 | "### Benchmark Keras prediction speed." 140 | ] 141 | }, 142 | { 143 | "metadata": { 144 | "id": "kGqN3UXquW-m", 145 | "colab_type": "code", 146 | "outputId": "450a9e18-310e-4ec7-a6ec-34f508e743c3", 147 | "colab": { 148 | "base_uri": "https://localhost:8080/", 149 | "height": 35 150 | } 151 | }, 152 | "cell_type": "code", 153 | "source": [ 154 | "import time\n", 155 | "times = []\n", 156 | "for i in range(20):\n", 157 | " start_time = time.time()\n", 158 | " preds = model.predict(x)\n", 159 | " delta = (time.time() - start_time)\n", 160 | " times.append(delta)\n", 161 | "mean_delta = np.array(times).mean()\n", 162 | "fps = 1/mean_delta\n", 163 | "print('average(sec):{},fps:{}'.format(mean_delta,fps))\n", 164 | "\n", 165 | "# Clear any previous session.\n", 166 | "tf.keras.backend.clear_session()" 167 | ], 168 | "execution_count": 4, 169 | "outputs": [ 170 | { 171 | "output_type": "stream", 172 | "text": [ 173 | "average(sec):0.016628634929656983,fps:60.137227393002135\n" 174 | ], 175 | "name": "stdout" 176 | } 177 | ] 178 | }, 179 | { 180 | "metadata": { 181 | "id": "FBsv7z2WukCz", 182 | "colab_type": "text" 183 | }, 184 | "cell_type": "markdown", 185 | "source": [ 186 | "## Freeze graph, generate `.pb` file.\n", 187 | "Take a notes of the input and output nodes names printed in the output, we will need them when converting `TensorRT` graph and prediction.\n", 188 | "\n", 189 | "For Keras MobileNetV2, they are,\n", 190 | "```\n", 191 | "['input_1'] ['Logits/Softmax']\n", 192 | "```\n" 193 | ] 194 | }, 195 | { 196 | "metadata": { 197 | "id": "L-Jx1Yq0uejv", 198 | "colab_type": "code", 199 | "outputId": "1f71e55c-8c33-4d66-92da-1240036c8bd8", 200 | "colab": { 201 | "base_uri": "https://localhost:8080/", 202 | "height": 274 203 | } 204 | }, 205 | "cell_type": "code", 206 | "source": [ 207 | "# force reset ipython namespaces\n", 208 | "%reset -f\n", 209 | "\n", 210 | "import tensorflow as tf\n", 211 | "from tensorflow.python.framework import graph_io\n", 212 | "from tensorflow.keras.models import load_model\n", 213 | "\n", 214 | "\n", 215 | "# Clear any previous session.\n", 216 | "tf.keras.backend.clear_session()\n", 217 | "\n", 218 | "save_pb_dir = './model'\n", 219 | "model_fname = './model/model.h5'\n", 220 | "def freeze_graph(graph, session, output, save_pb_dir='.', save_pb_name='frozen_model.pb', save_pb_as_text=False):\n", 221 | " with graph.as_default():\n", 222 | " graphdef_inf = tf.graph_util.remove_training_nodes(graph.as_graph_def())\n", 223 | " graphdef_frozen = tf.graph_util.convert_variables_to_constants(session, graphdef_inf, output)\n", 224 | " graph_io.write_graph(graphdef_frozen, save_pb_dir, save_pb_name, as_text=save_pb_as_text)\n", 225 | " return graphdef_frozen\n", 226 | "\n", 227 | "# This line must be executed before loading Keras model.\n", 228 | "tf.keras.backend.set_learning_phase(0) \n", 229 | "\n", 230 | "model = load_model(model_fname)\n", 231 | "\n", 232 | "session = tf.keras.backend.get_session()\n", 233 | "\n", 234 | "input_names = [t.op.name for t in model.inputs]\n", 235 | "output_names = [t.op.name for t in model.outputs]\n", 236 | "\n", 237 | "# Prints input and output nodes names, take notes of them.\n", 238 | "print(input_names, output_names)\n", 239 | "\n", 240 | "frozen_graph = freeze_graph(session.graph, session, [out.op.name for out in model.outputs], save_pb_dir=save_pb_dir)" 241 | ], 242 | "execution_count": 5, 243 | "outputs": [ 244 | { 245 | "output_type": "stream", 246 | "text": [ 247 | "WARNING:tensorflow:No training configuration found in save file: the model was *not* compiled. Compile it manually.\n", 248 | "['input_1'] ['Logits/Softmax']\n", 249 | "WARNING:tensorflow:From :15: remove_training_nodes (from tensorflow.python.framework.graph_util_impl) is deprecated and will be removed in a future version.\n", 250 | "Instructions for updating:\n", 251 | "Use tf.compat.v1.graph_util.remove_training_nodes\n", 252 | "WARNING:tensorflow:From :16: convert_variables_to_constants (from tensorflow.python.framework.graph_util_impl) is deprecated and will be removed in a future version.\n", 253 | "Instructions for updating:\n", 254 | "Use tf.compat.v1.graph_util.convert_variables_to_constants\n", 255 | "WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensorflow/python/framework/graph_util_impl.py:245: extract_sub_graph (from tensorflow.python.framework.graph_util_impl) is deprecated and will be removed in a future version.\n", 256 | "Instructions for updating:\n", 257 | "Use tf.compat.v1.graph_util.extract_sub_graph\n", 258 | "INFO:tensorflow:Froze 262 variables.\n", 259 | "INFO:tensorflow:Converted 262 variables to const ops.\n" 260 | ], 261 | "name": "stdout" 262 | } 263 | ] 264 | }, 265 | { 266 | "metadata": { 267 | "id": "38l-D915vahb", 268 | "colab_type": "text" 269 | }, 270 | "cell_type": "markdown", 271 | "source": [ 272 | "## Optimize with TensorRT\n", 273 | "Save the result as `.pb` file.\n", 274 | "\n", 275 | "*Notes: optimizing TensorRT graph can also be executed on Jetson Nano, but it is very slow.*" 276 | ] 277 | }, 278 | { 279 | "metadata": { 280 | "id": "g9qanB3Du2i-", 281 | "colab_type": "code", 282 | "outputId": "57f5f966-3b13-468d-de5e-fd812710027a", 283 | "colab": { 284 | "base_uri": "https://localhost:8080/", 285 | "height": 35 286 | } 287 | }, 288 | "cell_type": "code", 289 | "source": [ 290 | "import tensorflow.contrib.tensorrt as trt\n", 291 | "\n", 292 | "trt_graph = trt.create_inference_graph(\n", 293 | " input_graph_def=frozen_graph,\n", 294 | " outputs=output_names,\n", 295 | " max_batch_size=1,\n", 296 | " max_workspace_size_bytes=1 << 25,\n", 297 | " precision_mode='FP16',\n", 298 | " minimum_segment_size=50\n", 299 | ")" 300 | ], 301 | "execution_count": 7, 302 | "outputs": [ 303 | { 304 | "output_type": "stream", 305 | "text": [ 306 | "INFO:tensorflow:Running against TensorRT version 0.0.0\n" 307 | ], 308 | "name": "stdout" 309 | } 310 | ] 311 | }, 312 | { 313 | "metadata": { 314 | "id": "Azhh5OA2vI72", 315 | "colab_type": "code", 316 | "outputId": "0aa4643a-42cc-46de-df0a-97829cd9931a", 317 | "colab": { 318 | "base_uri": "https://localhost:8080/", 319 | "height": 35 320 | } 321 | }, 322 | "cell_type": "code", 323 | "source": [ 324 | "graph_io.write_graph(trt_graph, \"./model/\",\n", 325 | " \"trt_graph.pb\", as_text=False)" 326 | ], 327 | "execution_count": 8, 328 | "outputs": [ 329 | { 330 | "output_type": "execute_result", 331 | "data": { 332 | "text/plain": [ 333 | "'./model/trt_graph.pb'" 334 | ] 335 | }, 336 | "metadata": { 337 | "tags": [] 338 | }, 339 | "execution_count": 8 340 | } 341 | ] 342 | }, 343 | { 344 | "metadata": { 345 | "id": "NLg8O3dfvOum", 346 | "colab_type": "code", 347 | "outputId": "f7f6f202-e10f-4477-b035-a0f3149eda8b", 348 | "colab": { 349 | "base_uri": "https://localhost:8080/", 350 | "height": 126 351 | } 352 | }, 353 | "cell_type": "code", 354 | "source": [ 355 | "!ls model -alh" 356 | ], 357 | "execution_count": 9, 358 | "outputs": [ 359 | { 360 | "output_type": "stream", 361 | "text": [ 362 | "total 42M\n", 363 | "drwxr-xr-x 2 root root 4.0K Apr 13 05:02 .\n", 364 | "drwxr-xr-x 1 root root 4.0K Apr 13 05:01 ..\n", 365 | "-rw-r--r-- 1 root root 14M Apr 13 05:01 frozen_model.pb\n", 366 | "-rw-r--r-- 1 root root 14M Apr 13 05:01 model.h5\n", 367 | "-rw-r--r-- 1 root root 14M Apr 13 05:02 trt_graph.pb\n" 368 | ], 369 | "name": "stdout" 370 | } 371 | ] 372 | }, 373 | { 374 | "metadata": { 375 | "id": "jPVLp1FKvnSF", 376 | "colab_type": "text" 377 | }, 378 | "cell_type": "markdown", 379 | "source": [ 380 | "### Download the tensorRT graph `.pb` file from colab to your local machine." 381 | ] 382 | }, 383 | { 384 | "metadata": { 385 | "id": "ZrHyjN_Cvk4Z", 386 | "colab_type": "code", 387 | "colab": {} 388 | }, 389 | "cell_type": "code", 390 | "source": [ 391 | "from google.colab import files\n", 392 | "\n", 393 | "files.download('./model/trt_graph.pb')" 394 | ], 395 | "execution_count": 0, 396 | "outputs": [] 397 | }, 398 | { 399 | "metadata": { 400 | "id": "_Hg6qDGIwmQn", 401 | "colab_type": "text" 402 | }, 403 | "cell_type": "markdown", 404 | "source": [ 405 | "**Next step**: transfer the `trt_graph.pb` to your Jetson Nano, load it up and make predictions.\n", 406 | "\n", 407 | "\n", 408 | "`Step2_keras-jetson-ImageNet-predict.ipynb`" 409 | ] 410 | }, 411 | { 412 | "metadata": { 413 | "id": "4b6xtTKp_ve7", 414 | "colab_type": "code", 415 | "colab": {} 416 | }, 417 | "cell_type": "code", 418 | "source": [ 419 | "" 420 | ], 421 | "execution_count": 0, 422 | "outputs": [] 423 | } 424 | ] 425 | } -------------------------------------------------------------------------------- /Step1_Object_detection_Colab_TensorRT.ipynb: -------------------------------------------------------------------------------- 1 | { 2 | "nbformat": 4, 3 | "nbformat_minor": 0, 4 | "metadata": { 5 | "colab": { 6 | "name": "Step1_Object_detection_Colab_TensorRT.ipynb", 7 | "version": "0.3.2", 8 | "provenance": [], 9 | "collapsed_sections": [], 10 | "include_colab_link": true 11 | }, 12 | "kernelspec": { 13 | "name": "python3", 14 | "display_name": "Python 3" 15 | }, 16 | "accelerator": "GPU" 17 | }, 18 | "cells": [ 19 | { 20 | "cell_type": "markdown", 21 | "metadata": { 22 | "id": "view-in-github", 23 | "colab_type": "text" 24 | }, 25 | "source": [ 26 | "\"Open" 27 | ] 28 | }, 29 | { 30 | "metadata": { 31 | "id": "QQK7KlhRBfnf", 32 | "colab_type": "text" 33 | }, 34 | "cell_type": "markdown", 35 | "source": [ 36 | "# Step 1 : Convert TensorFlow object detection model into TensorRT model\n", 37 | "**For more detail, checkout [How to run TensorFlow Object Detection model on Jetson Nano](https://www.dlology.com/blog/how-to-run-tensorflow-object-detection-model-on-jetson-nano/) | DLology Blog**" 38 | ] 39 | }, 40 | { 41 | "metadata": { 42 | "id": "T_-qqpBAM3LW", 43 | "colab_type": "text" 44 | }, 45 | "cell_type": "markdown", 46 | "source": [ 47 | "## Pre-trained TensorFlow object detection models" 48 | ] 49 | }, 50 | { 51 | "metadata": { 52 | "id": "ucxso6LcMows", 53 | "colab_type": "code", 54 | "colab": {} 55 | }, 56 | "cell_type": "code", 57 | "source": [ 58 | "from collections import namedtuple\n", 59 | "DetectionModel = namedtuple('DetectionModel', ['name', 'url', 'extract_dir'])\n", 60 | "\n", 61 | "MODELS = {\n", 62 | " 'ssd_mobilenet_v1_coco': DetectionModel(\n", 63 | " 'ssd_mobilenet_v1_coco',\n", 64 | " 'http://download.tensorflow.org/models/object_detection/ssd_mobilenet_v1_coco_2018_01_28.tar.gz',\n", 65 | " 'ssd_mobilenet_v1_coco_2018_01_28',\n", 66 | " ),\n", 67 | " 'ssd_mobilenet_v2_coco': DetectionModel(\n", 68 | " 'ssd_mobilenet_v2_coco',\n", 69 | " 'http://download.tensorflow.org/models/object_detection/ssd_mobilenet_v2_coco_2018_03_29.tar.gz',\n", 70 | " 'ssd_mobilenet_v2_coco_2018_03_29',\n", 71 | " ),\n", 72 | " 'ssd_inception_v2_coco': DetectionModel(\n", 73 | " 'ssd_inception_v2_coco',\n", 74 | " 'http://download.tensorflow.org/models/object_detection/ssd_inception_v2_coco_2018_01_28.tar.gz',\n", 75 | " 'ssd_inception_v2_coco_2018_01_28',\n", 76 | " ),\n", 77 | " 'ssd_resnet_50_fpn_coco': DetectionModel(\n", 78 | " 'ssd_resnet_50_fpn_coco',\n", 79 | " 'http://download.tensorflow.org/models/object_detection/ssd_resnet50_v1_fpn_shared_box_predictor_640x640_coco14_sync_2018_07_03.tar.gz',\n", 80 | " 'ssd_resnet50_v1_fpn_shared_box_predictor_640x640_coco14_sync_2018_07_03',\n", 81 | " ),\n", 82 | " 'faster_rcnn_resnet50_coco': DetectionModel(\n", 83 | " 'faster_rcnn_resnet50_coco',\n", 84 | " 'http://download.tensorflow.org/models/object_detection/faster_rcnn_resnet50_coco_2018_01_28.tar.gz',\n", 85 | " 'faster_rcnn_resnet50_coco_2018_01_28',\n", 86 | " ),\n", 87 | " 'faster_rcnn_nas': DetectionModel(\n", 88 | " 'faster_rcnn_nas',\n", 89 | " 'http://download.tensorflow.org/models/object_detection/faster_rcnn_nas_coco_2018_01_28.tar.gz',\n", 90 | " 'faster_rcnn_nas_coco_2018_01_28',\n", 91 | " ),\n", 92 | " 'mask_rcnn_resnet50_atrous_coco': DetectionModel(\n", 93 | " 'mask_rcnn_resnet50_atrous_coco',\n", 94 | " 'http://download.tensorflow.org/models/object_detection/mask_rcnn_resnet50_atrous_coco_2018_01_28.tar.gz',\n", 95 | " 'mask_rcnn_resnet50_atrous_coco_2018_01_28',\n", 96 | " )\n", 97 | "}\n" 98 | ], 99 | "execution_count": 0, 100 | "outputs": [] 101 | }, 102 | { 103 | "metadata": { 104 | "id": "UyEDVcwfNDO-", 105 | "colab_type": "text" 106 | }, 107 | "cell_type": "markdown", 108 | "source": [ 109 | "## Select your model" 110 | ] 111 | }, 112 | { 113 | "metadata": { 114 | "id": "1sENfWFOIiTr", 115 | "colab_type": "code", 116 | "colab": {} 117 | }, 118 | "cell_type": "code", 119 | "source": [ 120 | "MODEL = 'ssd_mobilenet_v1_coco'" 121 | ], 122 | "execution_count": 0, 123 | "outputs": [] 124 | }, 125 | { 126 | "metadata": { 127 | "id": "MI6QhZndMZFV", 128 | "colab_type": "text" 129 | }, 130 | "cell_type": "markdown", 131 | "source": [ 132 | "## Install required packages" 133 | ] 134 | }, 135 | { 136 | "metadata": { 137 | "id": "GrcV6tNTDT8u", 138 | "colab_type": "code", 139 | "outputId": "f78f6b33-7004-4b0d-d1c1-7f61e79f2a63", 140 | "colab": { 141 | "base_uri": "https://localhost:8080/", 142 | "height": 290 143 | } 144 | }, 145 | "cell_type": "code", 146 | "source": [ 147 | "%cd /content\n", 148 | "!git clone --quiet https://github.com/tensorflow/models.git\n", 149 | "\n", 150 | "!apt-get install -qq protobuf-compiler python-pil python-lxml python-tk\n", 151 | "\n", 152 | "!pip install -q Cython contextlib2 pillow lxml matplotlib\n", 153 | "\n", 154 | "!pip install -q pycocotools\n", 155 | "\n", 156 | "%cd /content/models/research\n", 157 | "!protoc object_detection/protos/*.proto --python_out=.\n", 158 | "\n", 159 | "import os\n", 160 | "import sys\n", 161 | "os.environ['PYTHONPATH'] += ':/content/models/research/:/content/models/research/slim/'\n", 162 | "sys.path.append(\"/content/models/research/slim/\")\n", 163 | "\n", 164 | "!python object_detection/builders/model_builder_test.py" 165 | ], 166 | "execution_count": 0, 167 | "outputs": [ 168 | { 169 | "output_type": "stream", 170 | "text": [ 171 | "/content\n", 172 | "fatal: destination path 'models' already exists and is not an empty directory.\n", 173 | "/content/models/research\n", 174 | "\n", 175 | "WARNING: The TensorFlow contrib module will not be included in TensorFlow 2.0.\n", 176 | "For more information, please see:\n", 177 | " * https://github.com/tensorflow/community/blob/master/rfcs/20180907-contrib-sunset.md\n", 178 | " * https://github.com/tensorflow/addons\n", 179 | "If you depend on functionality not listed there, please file an issue.\n", 180 | "\n", 181 | "............s...\n", 182 | "----------------------------------------------------------------------\n", 183 | "Ran 16 tests in 0.062s\n", 184 | "\n", 185 | "OK (skipped=1)\n" 186 | ], 187 | "name": "stdout" 188 | } 189 | ] 190 | }, 191 | { 192 | "metadata": { 193 | "id": "jZVU5IGWLq87", 194 | "colab_type": "code", 195 | "colab": {} 196 | }, 197 | "cell_type": "code", 198 | "source": [ 199 | "import tensorflow.contrib.tensorrt as trt" 200 | ], 201 | "execution_count": 0, 202 | "outputs": [] 203 | }, 204 | { 205 | "metadata": { 206 | "id": "l9UHOzwHNsBd", 207 | "colab_type": "text" 208 | }, 209 | "cell_type": "markdown", 210 | "source": [ 211 | "`graph_utils.py`" 212 | ] 213 | }, 214 | { 215 | "metadata": { 216 | "id": "axmgTJ8BINlN", 217 | "colab_type": "code", 218 | "colab": {} 219 | }, 220 | "cell_type": "code", 221 | "source": [ 222 | "import tensorflow as tf\n", 223 | "\n", 224 | "\n", 225 | "def make_const6(const6_name='const6'):\n", 226 | " graph = tf.Graph()\n", 227 | " with graph.as_default():\n", 228 | " tf_6 = tf.constant(dtype=tf.float32, value=6.0, name=const6_name)\n", 229 | " return graph.as_graph_def()\n", 230 | "\n", 231 | "\n", 232 | "def make_relu6(output_name, input_name, const6_name='const6'):\n", 233 | " graph = tf.Graph()\n", 234 | " with graph.as_default():\n", 235 | " tf_x = tf.placeholder(tf.float32, [10, 10], name=input_name)\n", 236 | " tf_6 = tf.constant(dtype=tf.float32, value=6.0, name=const6_name)\n", 237 | " with tf.name_scope(output_name):\n", 238 | " tf_y1 = tf.nn.relu(tf_x, name='relu1')\n", 239 | " tf_y2 = tf.nn.relu(tf.subtract(tf_x, tf_6, name='sub1'), name='relu2')\n", 240 | "\n", 241 | " #tf_y = tf.nn.relu(tf.subtract(tf_6, tf.nn.relu(tf_x, name='relu1'), name='sub'), name='relu2')\n", 242 | " #tf_y = tf.subtract(tf_6, tf_y, name=output_name)\n", 243 | " tf_y = tf.subtract(tf_y1, tf_y2, name=output_name)\n", 244 | " \n", 245 | " graph_def = graph.as_graph_def()\n", 246 | " graph_def.node[-1].name = output_name\n", 247 | "\n", 248 | " # remove unused nodes\n", 249 | " for node in graph_def.node:\n", 250 | " if node.name == input_name:\n", 251 | " graph_def.node.remove(node)\n", 252 | " for node in graph_def.node:\n", 253 | " if node.name == const6_name:\n", 254 | " graph_def.node.remove(node)\n", 255 | " for node in graph_def.node:\n", 256 | " if node.op == '_Neg':\n", 257 | " node.op = 'Neg'\n", 258 | " \n", 259 | " return graph_def\n", 260 | "\n", 261 | "\n", 262 | "def convert_relu6(graph_def, const6_name='const6'):\n", 263 | " # add constant 6\n", 264 | " has_const6 = False\n", 265 | " for node in graph_def.node:\n", 266 | " if node.name == const6_name:\n", 267 | " has_const6 = True\n", 268 | " if not has_const6:\n", 269 | " const6_graph_def = make_const6(const6_name=const6_name)\n", 270 | " graph_def.node.extend(const6_graph_def.node)\n", 271 | " \n", 272 | " for node in graph_def.node:\n", 273 | " if node.op == 'Relu6':\n", 274 | " input_name = node.input[0]\n", 275 | " output_name = node.name\n", 276 | " relu6_graph_def = make_relu6(output_name, input_name, const6_name=const6_name)\n", 277 | " graph_def.node.remove(node)\n", 278 | " graph_def.node.extend(relu6_graph_def.node)\n", 279 | " \n", 280 | " return graph_def\n", 281 | "\n", 282 | "\n", 283 | "def remove_node(graph_def, node):\n", 284 | " for n in graph_def.node:\n", 285 | " if node.name in n.input:\n", 286 | " n.input.remove(node.name)\n", 287 | " ctrl_name = '^' + node.name\n", 288 | " if ctrl_name in n.input:\n", 289 | " n.input.remove(ctrl_name)\n", 290 | " graph_def.node.remove(node)\n", 291 | "\n", 292 | "\n", 293 | "def remove_op(graph_def, op_name):\n", 294 | " matches = [node for node in graph_def.node if node.op == op_name]\n", 295 | " for match in matches:\n", 296 | " remove_node(graph_def, match)\n", 297 | "\n", 298 | "\n", 299 | "def f_force_nms_cpu(frozen_graph):\n", 300 | " for node in frozen_graph.node:\n", 301 | " if 'NonMaxSuppression' in node.name:\n", 302 | " node.device = '/device:CPU:0'\n", 303 | " return frozen_graph\n", 304 | "\n", 305 | "\n", 306 | "def f_replace_relu6(frozen_graph):\n", 307 | " return convert_relu6(frozen_graph)\n", 308 | "\n", 309 | "\n", 310 | "def f_remove_assert(frozen_graph):\n", 311 | " remove_op(frozen_graph, 'Assert')\n", 312 | " return frozen_graph\n" 313 | ], 314 | "execution_count": 0, 315 | "outputs": [] 316 | }, 317 | { 318 | "metadata": { 319 | "id": "dm7-PKVDN1Ll", 320 | "colab_type": "text" 321 | }, 322 | "cell_type": "markdown", 323 | "source": [ 324 | "`detection.py`" 325 | ] 326 | }, 327 | { 328 | "metadata": { 329 | "id": "ihIdQpBFIKE4", 330 | "colab_type": "code", 331 | "colab": {} 332 | }, 333 | "cell_type": "code", 334 | "source": [ 335 | "from object_detection.protos import pipeline_pb2\n", 336 | "from object_detection import exporter\n", 337 | "\n", 338 | "import os\n", 339 | "import subprocess\n", 340 | "\n", 341 | "from google.protobuf import text_format\n", 342 | "\n", 343 | "import tensorflow as tf\n", 344 | "\n", 345 | "\n", 346 | "INPUT_NAME='image_tensor'\n", 347 | "BOXES_NAME='detection_boxes'\n", 348 | "CLASSES_NAME='detection_classes'\n", 349 | "SCORES_NAME='detection_scores'\n", 350 | "MASKS_NAME='detection_masks'\n", 351 | "NUM_DETECTIONS_NAME='num_detections'\n", 352 | "FROZEN_GRAPH_NAME='frozen_inference_graph.pb'\n", 353 | "PIPELINE_CONFIG_NAME='pipeline.config'\n", 354 | "CHECKPOINT_PREFIX='model.ckpt'\n", 355 | "\n", 356 | "\n", 357 | "\n", 358 | "def get_input_names(model):\n", 359 | " return [INPUT_NAME]\n", 360 | "\n", 361 | "\n", 362 | "def get_output_names(model):\n", 363 | " output_names = [BOXES_NAME, CLASSES_NAME, SCORES_NAME, NUM_DETECTIONS_NAME]\n", 364 | " if model == 'mask_rcnn_resnet50_atrous_coco':\n", 365 | " output_names.append(MASKS_NAME)\n", 366 | " return output_names\n", 367 | "\n", 368 | "\n", 369 | "def download_detection_model(model, output_dir='.'):\n", 370 | " \"\"\"Downloads a pre-trained object detection model\"\"\"\n", 371 | " global MODELS\n", 372 | "\n", 373 | " model_name = model\n", 374 | "\n", 375 | " model = MODELS[model_name]\n", 376 | " subprocess.call(['mkdir', '-p', output_dir])\n", 377 | " tar_file = os.path.join(output_dir, os.path.basename(model.url))\n", 378 | "\n", 379 | " config_path = os.path.join(output_dir, model.extract_dir, PIPELINE_CONFIG_NAME)\n", 380 | " checkpoint_path = os.path.join(output_dir, model.extract_dir, CHECKPOINT_PREFIX)\n", 381 | "\n", 382 | " if not os.path.exists(os.path.join(output_dir, model.extract_dir)):\n", 383 | " subprocess.call(['wget', model.url, '-O', tar_file])\n", 384 | " subprocess.call(['tar', '-xzf', tar_file, '-C', output_dir])\n", 385 | "\n", 386 | " # hack fix to handle mobilenet_v2 config bug\n", 387 | " subprocess.call(['sed', '-i', '/batch_norm_trainable/d', config_path])\n", 388 | "\n", 389 | " return config_path, checkpoint_path\n", 390 | "\n", 391 | "\n", 392 | "def build_detection_graph(config, checkpoint,\n", 393 | " batch_size=1,\n", 394 | " score_threshold=None,\n", 395 | " iou_threshold=None,\n", 396 | " force_nms_cpu=True,\n", 397 | " replace_relu6=True,\n", 398 | " remove_assert=True,\n", 399 | " input_shape=None,\n", 400 | " output_dir='.generated_model'):\n", 401 | " \"\"\"Builds a frozen graph for a pre-trained object detection model\"\"\"\n", 402 | " \n", 403 | " config_path = config\n", 404 | " checkpoint_path = checkpoint\n", 405 | "\n", 406 | " # parse config from file\n", 407 | " config = pipeline_pb2.TrainEvalPipelineConfig()\n", 408 | " with open(config_path, 'r') as f:\n", 409 | " text_format.Merge(f.read(), config, allow_unknown_extension=True)\n", 410 | "\n", 411 | " # override some config parameters\n", 412 | " if config.model.HasField('ssd'):\n", 413 | " config.model.ssd.feature_extractor.override_base_feature_extractor_hyperparams = True\n", 414 | " if score_threshold is not None:\n", 415 | " config.model.ssd.post_processing.batch_non_max_suppression.score_threshold = score_threshold\n", 416 | " if iou_threshold is not None:\n", 417 | " config.model.ssd.post_processing.batch_non_max_suppression.iou_threshold = iou_threshold\n", 418 | " if input_shape is not None:\n", 419 | " config.model.ssd.image_resizer.fixed_shape_resizer.height = input_shape[0]\n", 420 | " config.model.ssd.image_resizer.fixed_shape_resizer.width = input_shape[1]\n", 421 | " elif config.model.HasField('faster_rcnn'):\n", 422 | " if score_threshold is not None:\n", 423 | " config.model.faster_rcnn.second_stage_post_processing.score_threshold = score_threshold\n", 424 | " if input_shape is not None:\n", 425 | " config.model.faster_rcnn.image_resizer.fixed_shape_resizer.height = input_shape[0]\n", 426 | " config.model.faster_rcnn.image_resizer.fixed_shape_resizer.width = input_shape[1]\n", 427 | "\n", 428 | " if os.path.isdir(output_dir):\n", 429 | " subprocess.call(['rm', '-rf', output_dir])\n", 430 | "\n", 431 | " tf_config = tf.ConfigProto()\n", 432 | " tf_config.gpu_options.allow_growth = True\n", 433 | "\n", 434 | " # export inference graph to file (initial)\n", 435 | " with tf.Session(config=tf_config) as tf_sess:\n", 436 | " with tf.Graph().as_default() as tf_graph:\n", 437 | " exporter.export_inference_graph(\n", 438 | " 'image_tensor', \n", 439 | " config, \n", 440 | " checkpoint_path, \n", 441 | " output_dir, \n", 442 | " input_shape=[batch_size, None, None, 3]\n", 443 | " )\n", 444 | "\n", 445 | " # read frozen graph from file\n", 446 | " frozen_graph = tf.GraphDef()\n", 447 | " with open(os.path.join(output_dir, FROZEN_GRAPH_NAME), 'rb') as f:\n", 448 | " frozen_graph.ParseFromString(f.read())\n", 449 | "\n", 450 | " # apply graph modifications\n", 451 | " if force_nms_cpu:\n", 452 | " frozen_graph = f_force_nms_cpu(frozen_graph)\n", 453 | " if replace_relu6:\n", 454 | " frozen_graph = f_replace_relu6(frozen_graph)\n", 455 | " if remove_assert:\n", 456 | " frozen_graph = f_remove_assert(frozen_graph)\n", 457 | "\n", 458 | " # get input names\n", 459 | " # TODO: handle mask_rcnn \n", 460 | " input_names = [INPUT_NAME]\n", 461 | " output_names = [BOXES_NAME, CLASSES_NAME, SCORES_NAME, NUM_DETECTIONS_NAME]\n", 462 | "\n", 463 | " # remove temporary directory\n", 464 | " subprocess.call(['rm', '-rf', output_dir])\n", 465 | "\n", 466 | " return frozen_graph, input_names, output_names\n" 467 | ], 468 | "execution_count": 0, 469 | "outputs": [] 470 | }, 471 | { 472 | "metadata": { 473 | "id": "KLEiAFOAMNJ8", 474 | "colab_type": "text" 475 | }, 476 | "cell_type": "markdown", 477 | "source": [ 478 | "## Download pre-trained TensorFlow Object detection model" 479 | ] 480 | }, 481 | { 482 | "metadata": { 483 | "id": "f7WONX1MIp0k", 484 | "colab_type": "code", 485 | "colab": {} 486 | }, 487 | "cell_type": "code", 488 | "source": [ 489 | "config_path, checkpoint_path = download_detection_model(MODEL, 'data')" 490 | ], 491 | "execution_count": 0, 492 | "outputs": [] 493 | }, 494 | { 495 | "metadata": { 496 | "id": "fpj3sFQVMUFA", 497 | "colab_type": "text" 498 | }, 499 | "cell_type": "markdown", 500 | "source": [ 501 | "For improved performance, increase the non-max suppression score threshold in the downloaded config file from 1e-8 to something greater, like 0.1." 502 | ] 503 | }, 504 | { 505 | "metadata": { 506 | "id": "vdl8C-V8IrIA", 507 | "colab_type": "code", 508 | "outputId": "976b27c6-e8ee-4114-de52-94c09bf593b1", 509 | "colab": { 510 | "base_uri": "https://localhost:8080/", 511 | "height": 308 512 | } 513 | }, 514 | "cell_type": "code", 515 | "source": [ 516 | "frozen_graph, input_names, output_names = build_detection_graph(\n", 517 | " config=config_path,\n", 518 | " checkpoint=checkpoint_path,\n", 519 | " score_threshold=0.3,\n", 520 | " iou_threshold=0.5,\n", 521 | " batch_size=1\n", 522 | ")" 523 | ], 524 | "execution_count": 0, 525 | "outputs": [ 526 | { 527 | "output_type": "stream", 528 | "text": [ 529 | "INFO:tensorflow:depth of additional conv before box predictor: 0\n", 530 | "INFO:tensorflow:depth of additional conv before box predictor: 0\n", 531 | "INFO:tensorflow:depth of additional conv before box predictor: 0\n", 532 | "INFO:tensorflow:depth of additional conv before box predictor: 0\n", 533 | "INFO:tensorflow:depth of additional conv before box predictor: 0\n", 534 | "INFO:tensorflow:depth of additional conv before box predictor: 0\n" 535 | ], 536 | "name": "stdout" 537 | }, 538 | { 539 | "output_type": "stream", 540 | "text": [ 541 | "18 ops no flops stats due to incomplete shapes.\n", 542 | "18 ops no flops stats due to incomplete shapes.\n" 543 | ], 544 | "name": "stderr" 545 | }, 546 | { 547 | "output_type": "stream", 548 | "text": [ 549 | "INFO:tensorflow:Restoring parameters from data/ssd_mobilenet_v1_coco_2018_01_28/model.ckpt\n", 550 | "INFO:tensorflow:Restoring parameters from data/ssd_mobilenet_v1_coco_2018_01_28/model.ckpt\n", 551 | "INFO:tensorflow:Froze 199 variables.\n", 552 | "INFO:tensorflow:Converted 199 variables to const ops.\n", 553 | "INFO:tensorflow:No assets to save.\n", 554 | "INFO:tensorflow:No assets to write.\n", 555 | "INFO:tensorflow:SavedModel written to: .generated_model/saved_model/saved_model.pb\n", 556 | "INFO:tensorflow:Writing pipeline config file to .generated_model/pipeline.config\n" 557 | ], 558 | "name": "stdout" 559 | } 560 | ] 561 | }, 562 | { 563 | "metadata": { 564 | "id": "eWyu5YhLIrKe", 565 | "colab_type": "code", 566 | "outputId": "65331419-4fbe-4234-c6ee-4d781d7926b9", 567 | "colab": { 568 | "base_uri": "https://localhost:8080/", 569 | "height": 35 570 | } 571 | }, 572 | "cell_type": "code", 573 | "source": [ 574 | "print(output_names)" 575 | ], 576 | "execution_count": 0, 577 | "outputs": [ 578 | { 579 | "output_type": "stream", 580 | "text": [ 581 | "['detection_boxes', 'detection_classes', 'detection_scores', 'num_detections']\n" 582 | ], 583 | "name": "stdout" 584 | } 585 | ] 586 | }, 587 | { 588 | "metadata": { 589 | "id": "yBbfUCvEODTo", 590 | "colab_type": "text" 591 | }, 592 | "cell_type": "markdown", 593 | "source": [ 594 | "## Optimize the model with TensorRT" 595 | ] 596 | }, 597 | { 598 | "metadata": { 599 | "id": "uKC8ZWZEIwaS", 600 | "colab_type": "code", 601 | "outputId": "ba26eff1-0851-4306-e247-056f406fbd1f", 602 | "colab": { 603 | "base_uri": "https://localhost:8080/", 604 | "height": 35 605 | } 606 | }, 607 | "cell_type": "code", 608 | "source": [ 609 | "trt_graph = trt.create_inference_graph(\n", 610 | " input_graph_def=frozen_graph,\n", 611 | " outputs=output_names,\n", 612 | " max_batch_size=1,\n", 613 | " max_workspace_size_bytes=1 << 25,\n", 614 | " precision_mode='FP16',\n", 615 | " minimum_segment_size=50\n", 616 | ")" 617 | ], 618 | "execution_count": 0, 619 | "outputs": [ 620 | { 621 | "output_type": "stream", 622 | "text": [ 623 | "INFO:tensorflow:Running against TensorRT version 0.0.0\n" 624 | ], 625 | "name": "stdout" 626 | } 627 | ] 628 | }, 629 | { 630 | "metadata": { 631 | "id": "5vUsC1mCIwdz", 632 | "colab_type": "code", 633 | "colab": {} 634 | }, 635 | "cell_type": "code", 636 | "source": [ 637 | "with open('./data/trt_graph.pb', 'wb') as f:\n", 638 | " f.write(trt_graph.SerializeToString())" 639 | ], 640 | "execution_count": 0, 641 | "outputs": [] 642 | }, 643 | { 644 | "metadata": { 645 | "id": "KOgeJ-lyIwgU", 646 | "colab_type": "code", 647 | "outputId": "42bb3f5d-41e9-492c-fc07-052ddc4e0749", 648 | "colab": { 649 | "base_uri": "https://localhost:8080/", 650 | "height": 126 651 | } 652 | }, 653 | "cell_type": "code", 654 | "source": [ 655 | "!ls data -alh" 656 | ], 657 | "execution_count": 0, 658 | "outputs": [ 659 | { 660 | "output_type": "stream", 661 | "text": [ 662 | "total 101M\n", 663 | "drwxr-xr-x 3 root root 4.0K Apr 20 05:19 .\n", 664 | "drwxr-xr-x 72 root root 4.0K Apr 20 07:20 ..\n", 665 | "drwxr-xr-x 3 345018 5000 4.0K Apr 20 05:05 ssd_mobilenet_v1_coco_2018_01_28\n", 666 | "-rw-r--r-- 1 root root 73M Feb 10 2018 ssd_mobilenet_v1_coco_2018_01_28.tar.gz\n", 667 | "-rw-r--r-- 1 root root 28M Apr 20 07:20 trt_graph.pb\n" 668 | ], 669 | "name": "stdout" 670 | } 671 | ] 672 | }, 673 | { 674 | "metadata": { 675 | "id": "jWDYMW6ZAjbw", 676 | "colab_type": "text" 677 | }, 678 | "cell_type": "markdown", 679 | "source": [ 680 | "## Benchmark TensorFlow prediction speed" 681 | ] 682 | }, 683 | { 684 | "metadata": { 685 | "id": "RwUKKaM2_D95", 686 | "colab_type": "code", 687 | "colab": {} 688 | }, 689 | "cell_type": "code", 690 | "source": [ 691 | "input_names = ['image_tensor']\n", 692 | "output_names = ['detection_boxes', 'detection_classes', 'detection_scores', 'num_detections']\n", 693 | "# Create session and load graph\n", 694 | "tf_config = tf.ConfigProto()\n", 695 | "tf_config.gpu_options.allow_growth = True\n", 696 | "tf_sess = tf.Session(config=tf_config)\n", 697 | "tf.import_graph_def(frozen_graph, name='')\n", 698 | "\n", 699 | "tf_input = tf_sess.graph.get_tensor_by_name(input_names[0] + ':0')\n", 700 | "tf_scores = tf_sess.graph.get_tensor_by_name('detection_scores:0')\n", 701 | "tf_boxes = tf_sess.graph.get_tensor_by_name('detection_boxes:0')\n", 702 | "tf_classes = tf_sess.graph.get_tensor_by_name('detection_classes:0')\n", 703 | "tf_num_detections = tf_sess.graph.get_tensor_by_name('num_detections:0')" 704 | ], 705 | "execution_count": 0, 706 | "outputs": [] 707 | }, 708 | { 709 | "metadata": { 710 | "id": "ie0zwepg_MJA", 711 | "colab_type": "code", 712 | "colab": {} 713 | }, 714 | "cell_type": "code", 715 | "source": [ 716 | "import numpy as np\n", 717 | "image = np.random.random((300,300,3))\n", 718 | "scores, boxes, classes, num_detections = tf_sess.run([tf_scores, tf_boxes, tf_classes, tf_num_detections], feed_dict={\n", 719 | " tf_input: image[None, ...]\n", 720 | "})\n", 721 | "\n", 722 | "boxes = boxes[0] # index by 0 to remove batch dimension\n", 723 | "scores = scores[0]\n", 724 | "classes = classes[0]\n", 725 | "num_detections = num_detections[0]" 726 | ], 727 | "execution_count": 0, 728 | "outputs": [] 729 | }, 730 | { 731 | "metadata": { 732 | "id": "1e6NHFOWAREp", 733 | "colab_type": "code", 734 | "outputId": "f4ebf3b0-6e39-4f8d-b497-57ae82af0186", 735 | "colab": { 736 | "base_uri": "https://localhost:8080/", 737 | "height": 35 738 | } 739 | }, 740 | "cell_type": "code", 741 | "source": [ 742 | "import time\n", 743 | "times = []\n", 744 | "for i in range(20):\n", 745 | " start_time = time.time()\n", 746 | " scores, boxes, classes, num_detections = tf_sess.run([tf_scores, tf_boxes, tf_classes, tf_num_detections], feed_dict={\n", 747 | " tf_input: image[None, ...]\n", 748 | " })\n", 749 | "\n", 750 | " delta = (time.time() - start_time)\n", 751 | " times.append(delta)\n", 752 | "mean_delta = np.array(times).mean()\n", 753 | "fps = 1/mean_delta\n", 754 | "print('average(sec):{},fps:{}'.format(mean_delta,fps))\n" 755 | ], 756 | "execution_count": 0, 757 | "outputs": [ 758 | { 759 | "output_type": "stream", 760 | "text": [ 761 | "average(sec):0.010851287841796875,fps:92.15496027560994\n" 762 | ], 763 | "name": "stdout" 764 | } 765 | ] 766 | }, 767 | { 768 | "metadata": { 769 | "id": "ySxWp6zQBFL7", 770 | "colab_type": "code", 771 | "colab": {} 772 | }, 773 | "cell_type": "code", 774 | "source": [ 775 | "tf_sess.close()" 776 | ], 777 | "execution_count": 0, 778 | "outputs": [] 779 | }, 780 | { 781 | "metadata": { 782 | "id": "jPVLp1FKvnSF", 783 | "colab_type": "text" 784 | }, 785 | "cell_type": "markdown", 786 | "source": [ 787 | "### Download the tensorRT graph `.pb` file from colab to your local machine." 788 | ] 789 | }, 790 | { 791 | "metadata": { 792 | "id": "ZrHyjN_Cvk4Z", 793 | "colab_type": "code", 794 | "colab": {} 795 | }, 796 | "cell_type": "code", 797 | "source": [ 798 | "from google.colab import files\n", 799 | "\n", 800 | "files.download('./data/trt_graph.pb')" 801 | ], 802 | "execution_count": 0, 803 | "outputs": [] 804 | }, 805 | { 806 | "metadata": { 807 | "id": "_Hg6qDGIwmQn", 808 | "colab_type": "text" 809 | }, 810 | "cell_type": "markdown", 811 | "source": [ 812 | "**Next step**: transfer the `trt_graph.pb` to your Jetson Nano, load it up and make predictions.\n", 813 | "\n", 814 | "\n", 815 | "`Step2_jetson-object-detection-predict.ipynb`" 816 | ] 817 | }, 818 | { 819 | "metadata": { 820 | "id": "UkT7r-rnNj5j", 821 | "colab_type": "code", 822 | "colab": {} 823 | }, 824 | "cell_type": "code", 825 | "source": [ 826 | "" 827 | ], 828 | "execution_count": 0, 829 | "outputs": [] 830 | } 831 | ] 832 | } --------------------------------------------------------------------------------