{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# Using TensorRT to Optimize Caffe Models in Python\n", "\n", "TensorRT 4.0 includes support for a Python API to load in and optimize Caffe models, which can then be executed and stored.\n", "\n", "First, we import TensorRT." ] }, { "cell_type": "code", "execution_count": 1, "metadata": {}, "outputs": [], "source": [ "import tensorrt as trt" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "We use PyCUDA to transfer data to/from the GPU and NumPy to store data." ] }, { "cell_type": "code", "execution_count": 2, "metadata": {}, "outputs": [], "source": [ "import pycuda.driver as cuda\n", "import pycuda.autoinit\n", "import numpy as np" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "For this example, we also want to import `pillow` (an image processing library) and `randint`." ] }, { "cell_type": "code", "execution_count": 3, "metadata": {}, "outputs": [], "source": [ "from random import randint\n", "from PIL import Image\n", "from matplotlib.pyplot import imshow #to show test case" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Since we are converting a Caffe model, we also need to use the `caffeparser` located in `tensorrt.parsers`." ] }, { "cell_type": "code", "execution_count": 4, "metadata": {}, "outputs": [], "source": [ "from tensorrt import parsers" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Typically, the first thing you will do is create a logger, which is used extensively during the model conversion and inference process. We provide a simple logger implementation in `tensorrt.infer.ConsoleLogger`, but you can define your own as well." ] }, { "cell_type": "code", "execution_count": 5, "metadata": {}, "outputs": [], "source": [ "G_LOGGER = trt.infer.ConsoleLogger(trt.infer.LogSeverity.ERROR)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Now, we define some constants for our model, which we will use to classify digits from the MNIST dataset." ] }, { "cell_type": "code", "execution_count": 6, "metadata": {}, "outputs": [], "source": [ "INPUT_LAYERS = ['data']\n", "OUTPUT_LAYERS = ['prob']\n", "INPUT_H = 28\n", "INPUT_W = 28\n", "OUTPUT_SIZE = 10" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "We now define data paths. Please change these to reflect where you placed the data included with the samples." ] }, { "cell_type": "code", "execution_count": 7, "metadata": {}, "outputs": [], "source": [ "MODEL_PROTOTXT = './data/mnist/mnist.prototxt'\n", "CAFFE_MODEL = './data/mnist/mnist.caffemodel'\n", "DATA = './data/mnist/'\n", "IMAGE_MEAN = './data/mnist/mnist_mean.binaryproto'" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "The first step is to create our engine. The Python API provides utilities to make this simpler. Here, we use the caffe model converter utility in `tensorrt.utils`. We provide a logger, a path to the model prototxt, the model file, the max batch size, the max workspace size, the output layer(s) and the data type of the weights." ] }, { "cell_type": "code", "execution_count": 8, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Building Engine\n" ] } ], "source": [ "engine = trt.utils.caffe_to_trt_engine(G_LOGGER,\n", " MODEL_PROTOTXT,\n", " CAFFE_MODEL,\n", " 1, \n", " 1 << 20, \n", " OUTPUT_LAYERS,\n", " trt.infer.DataType.FLOAT)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Now let's generate a test case for our engine." ] }, { "cell_type": "code", "execution_count": 9, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Test Case: 4\n" ] }, { "data": { "image/png": "iVBORw0KGgoAAAANSUhEUgAAAP8AAAD8CAYAAAC4nHJkAAAABHNCSVQICAgIfAhkiAAAAAlwSFlzAAALEgAACxIB0t1+/AAAADl0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uIDIuMi4yLCBodHRwOi8vbWF0cGxvdGxpYi5vcmcvhp/UCwAADJNJREFUeJzt3WuMXHUdxvHnYSmJtmhoKZu1bClgxRCIlayNRqIYb7RqikaJNTE1Ia6J1kuiiYgvxBcmaATCC6NZpFqMgiZA6ItGxUZTr8BCyrUqiIu0lHZbiALKpdufL+ag27JzZjpzZs5sf99Pspkz53/OnF9O++y5/M/O3xEhAPkcV3cBAOpB+IGkCD+QFOEHkiL8QFKEH0iK8ANJEX4gKcIPJHV8Pzd28uKhWDG6oJ+bBFKZeuxF7X9yxu0s21X4bV8o6RpJQ5K+HxFXlC2/YnSB7vjFaDebBFBi9Xsfa3vZjk/7bQ9J+o6kNZLOlrTe9tmdfh6A/urmmn+1pIcj4pGIeEHSjZLWVVMWgF7rJvzLJM0+x9hVzDuM7XHbk7Ynpw/MdLE5AFXq+d3+iJiIiLGIGFu6ZKjXmwPQpm7Cv1vS7Lt3pxbzAMwD3YT/TkkrbZ9u+wRJH5W0pZqyAPRax119EXHQ9kZJv1Cjq29TRDxQWWUAeqqrfv6I2Cppa0W1AOgjHu8FkiL8QFKEH0iK8ANJEX4gKcIPJEX4gaQIP5AU4QeSIvxAUoQfSIrwA0kRfiApwg8kRfiBpAg/kBThB5Ii/EBShB9IivADSRF+ICnCDyRF+IGkCD+QFOEHkiL8QFKEH0iK8ANJEX4gqa5G6bU9JelpSTOSDkbEWBVF4dixZu3HmrY9+83/lK67/dxbqi4Hs3QV/sI7ImJ/BZ8DoI847QeS6jb8IemXtu+yPV5FQQD6o9vT/vMjYrftUyTdZvvPEbF99gLFL4VxSVq+rIqrDABV6OrIHxG7i9d9km6RtHqOZSYiYiwixpYuGepmcwAq1HH4bS+0feJL05LeI+n+qgoD0FvdnIcPS7rF9kuf85OI+HklVQHouY7DHxGPSHpDhbXgGBSNg8Ocpv8wUr7yuRUXg8PQ1QckRfiBpAg/kBThB5Ii/EBShB9IiudtjwE7nn++aduHtn6udN3tH7iytP3U4xd1VBMGH0d+ICnCDyRF+IGkCD+QFOEHkiL8QFKEH0iKfv5jwI7nR5u2rdx4R+m6nz7nw6XtW1byFQ3HKo78QFKEH0iK8ANJEX4gKcIPJEX4gaQIP5AU4QeSIvxAUoQfSIrwA0kRfiApwg8kRfiBpAg/kFTLv+e3vUnS+yXti4hzinmLJf1U0gpJU5Iujoinelcmynzj7rVN287QPaXr/uWJU8o/fGUnFf3fv5cvbNq28PHo7sPRlXaO/D+UdOER8y6VtC0iVkraVrwHMI+0DH9EbJf05BGz10naXExvlnRRxXUB6LFOr/mHI2JPMf2EpOGK6gHQJ13f8IuIkNT04s32uO1J25PTB2a63RyAinQa/r22RySpeN3XbMGImIiIsYgYW7pkqMPNAahap+HfImlDMb1B0q3VlAOgX1qG3/YNkv4o6Szbu2xfIukKSe+2/ZCkdxXvAcwjLfv5I2J9k6Z3VlwLOnTc317R8bqvWfzPCit5ud1vb358ef1V/yhf+esVF4PD8IQfkBThB5Ii/EBShB9IivADSRF+ICmG6J4H9s88W9p+5o1H/t3V/x06fkHpumeceKCjmjD/ceQHkiL8QFKEH0iK8ANJEX4gKcIPJEX4gaTo558H3nLDl0rbz3jwT03bfMIJpev+9jfnlrafPnxWaXsrr7ux+TMKj190Wlefje5w5AeSIvxAUoQfSIrwA0kRfiApwg8kRfiBpOjnnwded82jpe0HS9rihRdK1z39K82fEWh8QIthtO3y9hLPrXlLx+uiexz5gaQIP5AU4QeSIvxAUoQfSIrwA0kRfiCplv38tjdJer+kfRFxTjHvckmflDRdLHZZRGztVZHpHTrU8arPrxkrbX9u41Ol7ect3dXxtiXp0Y+c0ryxxSME6K12jvw/lHThHPOvjohVxQ/BB+aZluGPiO2Smg8JA2Be6uaaf6Pte21vsn1SZRUB6ItOw/9dSWdKWiVpj6Qrmy1oe9z2pO3J6QMzHW4OQNU6Cn9E7I2ImYg4JOlaSatLlp2IiLGIGFu6ZKjTOgFUrKPw2x6Z9faDku6vphwA/dJOV98Nki6QdLLtXZK+JukC26vU6KyZkvSpHtYIoAdahj8i1s8x+7oe1IImnnrbitL2E8eb/zNuPes7peu+8rjy7/Vv5aZnXlXavum5Vzdv7PyrAFABnvADkiL8QFKEH0iK8ANJEX4gKcIPJMVXd88Df7j6e12s3V1XXivPHir//DhY9sXiqBNHfiApwg8kRfiBpAg/kBThB5Ii/EBShB9Iin5+dOWeZ5eXts8c4LtfBxVHfiApwg8kRfiBpAg/kBThB5Ii/EBShB9IivADSRF+ICnCDyRF+IGkCD+QFOEHkiL8QFKEH0iq5d/z2x6VdL2kYUkhaSIirrG9WNJPJa2QNCXp4oh4qnelYhDduf+00vZXaKpp26I37a+4GhyNdo78ByV9MSLOlvRmSZ+xfbakSyVti4iVkrYV7wHMEy3DHxF7IuLuYvppSTslLZO0TtLmYrHNki7qVZEAqndU1/y2V0h6o6TbJQ1HxJ6i6Qk1LgsAzBNth9/2Ikk3SfpCRPxrdltEhBr3A+Zab9z2pO3J6QMzXRULoDpthd/2AjWC/+OIuLmYvdf2SNE+ImnfXOtGxEREjEXE2NIlQ1XUDKACLcNv25Kuk7QzIq6a1bRF0oZieoOkW6svD0CvtPPV3W+V9HFJ99neUcy7TNIVkn5m+xJJj0q6uDclYpBN/36ktH15SVffBa95qOJqcDRahj8ififJTZrfWW05APqFJ/yApAg/kBThB5Ii/EBShB9IivADSRF+ICnCDyRF+IGkCD+QFOEHkiL8QFKEH0iK8ANJtfP3/EDHhl7/2qZtnz35By3WXlRtMTgMR34gKcIPJEX4gaQIP5AU4QeSIvxAUoQfSIp+fnTludEXS9t3fvakpm3Lj6cfv04c+YGkCD+QFOEHkiL8QFKEH0iK8ANJEX4gqZb9/LZHJV0vaVhSSJqIiGtsXy7pk5Kmi0Uvi4itvSoUg+nv77u27hLQoXYe8jko6YsRcbftEyXdZfu2ou3qiPh278oD0Cstwx8ReyTtKaaftr1T0rJeFwagt47qmt/2CklvlHR7MWuj7Xttb7I953OctsdtT9qenD4w01WxAKrTdvhtL5J0k6QvRMS/JH1X0pmSVqlxZnDlXOtFxEREjEXE2NIlQxWUDKAKbYXf9gI1gv/jiLhZkiJib0TMRMQhSddKWt27MgFUrWX4bVvSdZJ2RsRVs+aPzFrsg5Lur748AL3Szt3+t0r6uKT7bO8o5l0mab3tVWp0/01J+lRPKgTQE+3c7f+dJM/RRJ8+MI/xhB+QFOEHkiL8QFKEH0iK8ANJEX4gKcIPJEX4gaQIP5AU4QeSIvxAUoQfSIrwA0kRfiApR0T/NmZPS3p01qyTJe3vWwFHZ1BrG9S6JGrrVJW1nRYRS9tZsK/hf9nG7cmIGKutgBKDWtug1iVRW6fqqo3TfiApwg8kVXf4J2refplBrW1Q65KorVO11FbrNT+A+tR95AdQk1rCb/tC23+x/bDtS+uooRnbU7bvs73D9mTNtWyyvc/2/bPmLbZ9m+2Hitc5h0mrqbbLbe8u9t0O22trqm3U9q9tP2j7AdufL+bXuu9K6qplv/X9tN/2kKS/Snq3pF2S7pS0PiIe7GshTdiekjQWEbX3Cdt+m6RnJF0fEecU874l6cmIuKL4xXlSRHx5QGq7XNIzdY/cXAwoMzJ7ZGlJF0n6hGrcdyV1Xawa9lsdR/7Vkh6OiEci4gVJN0paV0MdAy8itkt68ojZ6yRtLqY3q/Gfp++a1DYQImJPRNxdTD8t6aWRpWvddyV11aKO8C+T9Nis97s0WEN+h6Rf2r7L9njdxcxhuBg2XZKekDRcZzFzaDlycz8dMbL0wOy7Tka8rho3/F7u/Ig4T9IaSZ8pTm8HUjSu2Qapu6atkZv7ZY6Rpf+nzn3X6YjXVasj/Lsljc56f2oxbyBExO7idZ+kWzR4ow/vfWmQ1OJ1X831/M8gjdw818jSGoB9N0gjXtcR/jslrbR9uu0TJH1U0pYa6ngZ2wuLGzGyvVDSezR4ow9vkbShmN4g6dYaaznMoIzc3GxkadW87wZuxOuI6PuPpLVq3PH/m6Sv1lFDk7rOkHRP8fNA3bVJukGN08AX1bg3comkJZK2SXpI0q8kLR6g2n4k6T5J96oRtJGaajtfjVP6eyXtKH7W1r3vSuqqZb/xhB+QFDf8gKQIP5AU4QeSIvxAUoQfSIrwA0kRfiApwg8k9V/YvtIY2G31ZAAAAABJRU5ErkJggg==\n", "text/plain": [ "
" ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "rand_file = randint(0,9)\n", "path = DATA + str(rand_file) + '.pgm'\n", "im = Image.open(path)\n", "%matplotlib inline\n", "imshow(np.asarray(im))\n", "arr = np.array(im)\n", "img = arr.ravel()\n", "print(\"Test Case: \" + str(rand_file))" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "We now need to apply the mean to the input image; we have this stored in a .binaryproto file which can be read by the caffeparser." ] }, { "cell_type": "code", "execution_count": 10, "metadata": {}, "outputs": [], "source": [ "parser = parsers.caffeparser.create_caffe_parser()\n", "mean_blob = parser.parse_binary_proto(IMAGE_MEAN)\n", "parser.destroy()\n", "#NOTE: This is different than the C++ API, you must provide the size of the data\n", "mean = mean_blob.get_data(INPUT_W ** 2) \n", "data = np.empty([INPUT_W ** 2])\n", "for i in range(INPUT_W ** 2):\n", " data[i] = float(img[i]) - mean[i]\n", "mean_blob.destroy() " ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Now we need to create a runtime for inference and create a context for our engine." ] }, { "cell_type": "code", "execution_count": 11, "metadata": {}, "outputs": [], "source": [ "runtime = trt.infer.create_infer_runtime(G_LOGGER)\n", "context = engine.create_execution_context()" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Now we can run inference. We first make sure our data is in the correct datatype (FP32 for this model). Then, we allocate an empty array on the CPU to store our results from inference." ] }, { "cell_type": "code", "execution_count": 12, "metadata": {}, "outputs": [], "source": [ "assert(engine.get_nb_bindings() == 2)\n", "#convert input data to Float32\n", "img = img.astype(np.float32)\n", "#create output array to receive data \n", "output = np.empty(OUTPUT_SIZE, dtype = np.float32)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Now, we allocate memory on the GPU with PyCUDA and register it with the engine. The size of the allocations is the size of the input/expected output * the batch size." ] }, { "cell_type": "code", "execution_count": 13, "metadata": {}, "outputs": [], "source": [ "d_input = cuda.mem_alloc(1 * img.size * img.dtype.itemsize)\n", "d_output = cuda.mem_alloc(1 * output.size * output.dtype.itemsize)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "The engine needs bindings provided as pointers to the GPU memory. PyCUDA lets us do this for memory allocations by casting those allocations to ints." ] }, { "cell_type": "code", "execution_count": 14, "metadata": {}, "outputs": [], "source": [ "bindings = [int(d_input), int(d_output)]" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "We then create a cuda stream to run inference." ] }, { "cell_type": "code", "execution_count": 15, "metadata": {}, "outputs": [], "source": [ "stream = cuda.Stream()" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Now we transfer the data to the GPU, run inference and the copy the results back." ] }, { "cell_type": "code", "execution_count": 16, "metadata": {}, "outputs": [], "source": [ "#transfer input data to device\n", "cuda.memcpy_htod_async(d_input, img, stream)\n", "#execute model \n", "context.enqueue(1, bindings, stream.handle, None)\n", "#transfer predictions back\n", "cuda.memcpy_dtoh_async(output, d_output, stream)\n", "#syncronize threads\n", "stream.synchronize()" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "We can run `np.argmax` to get a prediction." ] }, { "cell_type": "code", "execution_count": 17, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Test Case: 4\n", "Prediction: 4\n" ] } ], "source": [ "print(\"Test Case: \" + str(rand_file))\n", "print (\"Prediction: \" + str(np.argmax(output)))" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "We can also write our engine to a file to use later." ] }, { "cell_type": "code", "execution_count": 18, "metadata": {}, "outputs": [ { "data": { "text/plain": [ "True" ] }, "execution_count": 18, "metadata": {}, "output_type": "execute_result" } ], "source": [ "trt.utils.write_engine_to_file(\"./data/mnist/new_mnist.engine\", engine.serialize())" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "You can load this engine later by using `tensorrt.utils.load_engine`." ] }, { "cell_type": "code", "execution_count": 19, "metadata": {}, "outputs": [], "source": [ "new_engine = trt.utils.load_engine(G_LOGGER, \"./data/mnist/new_mnist.engine\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Finally, we clean up our context, engine and runtime" ] }, { "cell_type": "code", "execution_count": 20, "metadata": {}, "outputs": [], "source": [ "context.destroy()\n", "engine.destroy()\n", "new_engine.destroy()\n", "runtime.destroy()" ] } ], "metadata": { "kernelspec": { "display_name": "Python 3", "language": "python", "name": "python3" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 2 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython2", "version": "2.7.12" } }, "nbformat": 4, "nbformat_minor": 2 }