<jupyter_start><jupyter_text>Using TensorRT to Optimize Caffe Models in PythonTensorRT 3.0.4 includes support for a Python API to load in and optimize Caffe models which can then be executed and stored as portable PLAN files. The following notebook explains the workflow you can use to do so. For python the TensorRT library is refered to as ```tensorrt```, for the Early Access you should have been provided a wheel file with the API, this can be installed by using ```pip``` (e.g. for python2.7 on ubuntu 16.04 - ```pip install tensorrt-3.0.4-cp27-cp27mu-linux_x86_64.whl```). You can import tensorrt as you would import any other package<jupyter_code>import tensorrt as trt<jupyter_output><empty_output><jupyter_text>There are also some common tools that are used with tensorrt typically. We use PyCUDA to handle the CUDA operations needed to allocate memory on your GPU and to transfer data to the GPU and results back to the CPU. We also use numpy as our primary method to store data<jupyter_code>import pycuda.driver as cuda
import pycuda.autoinit
import numpy as np<jupyter_output><empty_output><jupyter_text>For this example we also want to import an image processing library (```pillow``` in this case) and ```randint```<jupyter_code>from random import randint
from PIL import Image
from matplotlib.pyplot import imshow #to show test case<jupyter_output><empty_output><jupyter_text>Since we are converting a Caffe model we also need to use the ```caffeparser``` which is located in ```tensorrt.parsers```<jupyter_code>from tensorrt import parsers<jupyter_output><empty_output><jupyter_text>Typically the first thing you will do is create a logger, which is used in may places during the model conversion and inference process. We provide a simple logger implementation in ```tensorrt.infer.ConsoleLogger``` but in the RC you will be able to define your own as well.<jupyter_code>G_LOGGER = trt.infer.ConsoleLogger(trt.infer.LogSeverity.ERROR)<jupyter_output><empty_output><jupyter_text>Now we will define some constants about our model, which in this example we will used to classify digits from the MNIST dataset<jupyter_code>INPUT_LAYERS = ['data']
OUTPUT_LAYERS = ['prob']
INPUT_H = 28
INPUT_W =  28
OUTPUT_SIZE = 10<jupyter_output><empty_output><jupyter_text>We are also going to define some paths, please change these to reflect where you placed the data included with the samples<jupyter_code>MODEL_PROTOTXT = './data/mnist/mnist.prototxt'
CAFFE_MODEL = './data/mnist/mnist.caffemodel'
DATA = './data/mnist/'
IMAGE_MEAN = './data/mnist/mnist_mean.binaryproto'<jupyter_output><empty_output><jupyter_text>The first step is to create our engine. The Python API provides some nice utilities to make this much simplier. Here we use the caffe model converter utility in ```tensorrt.utils```. We provide it 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.<jupyter_code>engine = trt.utils.caffe_to_trt_engine(G_LOGGER,
                                       MODEL_PROTOTXT,
                                       CAFFE_MODEL,
                                       1, 
                                       1 << 20, 
                                       OUTPUT_LAYERS,
                                       trt.infer.DataType.FLOAT)<jupyter_output><empty_output><jupyter_text>Now let's generate a test case for our engine.<jupyter_code>rand_file = randint(0,9)
path = DATA + str(rand_file) + '.pgm'
im = Image.open(path)
%matplotlib inline
imshow(np.asarray(im))
arr = np.array(im)
img = arr.ravel()
print("Test Case: " + str(rand_file))<jupyter_output><empty_output><jupyter_text>We now need to apply the mean to the input image, we have this stored in a .binaryproto file which we use the caffeparser to read<jupyter_code>parser = parsers.caffeparser.create_caffe_parser()
mean_blob = parser.parse_binary_proto(IMAGE_MEAN)
parser.destroy()
#NOTE: This is different than the C++ API, you must provide the size of the data
mean = mean_blob.get_data(INPUT_W ** 2) 
data = np.empty([INPUT_W ** 2])
for i in range(INPUT_W ** 2):
    data[i] = float(img[i]) - mean[i]
mean_blob.destroy()<jupyter_output><empty_output><jupyter_text>Now we need to create a runtime for Inference and create a context for our engine<jupyter_code>runtime = trt.infer.create_infer_runtime(G_LOGGER)
context = engine.create_execution_context()<jupyter_output><empty_output><jupyter_text>Now we can run inference, we are going to start by making sure our data is in the correct datatype (FP32 for this model). Then we are going to create an empty array on the CPU to hold our results from inference.<jupyter_code>assert(engine.get_nb_bindings() == 2)
#convert input data to Float32
img = img.astype(np.float32)
#create output array to receive data 
output = np.empty(OUTPUT_SIZE, dtype = np.float32)<jupyter_output><empty_output><jupyter_text>Now we are going to allocate memory on the GPU with PyCUDA and register them with the engine. The size of the allocations is the size of the input and expected output * the batch size<jupyter_code>d_input = cuda.mem_alloc(1 * img.size * img.dtype.itemsize)
d_output = cuda.mem_alloc(1 * output.size * output.dtype.itemsize)<jupyter_output><empty_output><jupyter_text>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<jupyter_code>bindings = [int(d_input), int(d_output)]<jupyter_output><empty_output><jupyter_text>We also are going to create a cuda stream to run inference in<jupyter_code>stream = cuda.Stream()<jupyter_output><empty_output><jupyter_text>Now we are going to transfer the data to the GPU, run inference and the copy the results back.<jupyter_code>#transfer input data to device
cuda.memcpy_htod_async(d_input, img, stream)
#execute model 
context.enqueue(1, bindings, stream.handle, None)
#transfer predictions back
cuda.memcpy_dtoh_async(output, d_output, stream)
#syncronize threads
stream.synchronize()<jupyter_output><empty_output><jupyter_text>Now we have our results. We can just run ArgMax to get a prediction<jupyter_code>print("Test Case: " + str(rand_file))
print ("Prediction: " + str(np.argmax(output)))<jupyter_output><empty_output><jupyter_text>We can also save our engine to a file to use later<jupyter_code>trt.utils.write_engine_to_file("./data/mnist/new_mnist.engine", engine.serialize())<jupyter_output><empty_output><jupyter_text>You can then load this engine later by using ```tensorrt.utils.load_engine```<jupyter_code>new_engine = trt.utils.load_engine(G_LOGGER, "./data/mnist/new_mnist.engine")<jupyter_output><empty_output><jupyter_text>And as a final step, we are going to clean up our context, engine and runtime<jupyter_code>context.destroy()
engine.destroy()
new_engine.destroy()
runtime.destroy()<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.08519269776876268, 'fraction_numerical': 0.006809620399884092, 'mean_word_length': 5.0130662020905925, 'pattern_counts': {'":': 0, '<': 83, '<?xml version=': 0, '>': 81, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 7, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '20', 'dir': 'jupyter-structured-clean-dedup', 'id': '364857', 'license': 'MIT', 'n_tokens_mistral': 2204, 'n_tokens_neox': 2128, 'n_words': 727, 'path': 'src/tensorrt/tensorrt-3.0.4/caffe_to_tensorrt.ipynb', 'repo_name': 'aimuch/AIplay', 'size': '6902'}