From 3c5939cd0d51bb86b2da0f6fe0f11df8a1fdef16 Mon Sep 17 00:00:00 2001 From: irvingzhang0512 Date: Wed, 14 Apr 2021 15:34:57 +0800 Subject: [PATCH] add tensorrt python api sample for lenet (#476) --- lenet/README.md | 18 +++++ lenet/lenet.py | 190 ++++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 208 insertions(+) create mode 100644 lenet/lenet.py diff --git a/lenet/README.md b/lenet/README.md index 0560598..c2656dc 100644 --- a/lenet/README.md +++ b/lenet/README.md @@ -2,6 +2,8 @@ lenet5 is the simplest net in this tensorrtx project. You can learn the basic procedures of building tensorrt app from API. Including `define network`, `build engine`, `set output`, `do inference`, `serialize model to file`, `deserialize model from file`, etc. +## TensorRT C++ API + ``` // 1. generate lenet5.wts from https://github.com/wang-xinyu/pytorchx/tree/master/lenet @@ -26,4 +28,20 @@ sudo ./lenet -d // deserialize plan file and run inference // 4. see if the output is same as pytorchx/lenet ``` +## TensorRT Python API +``` +# 1. generate lenet5.wts from https://github.com/wang-xinyu/pytorchx/tree/master/lenet + +# 2. put lenet5.wts into tensorrtx/lenet + +# 3. install Python dependencies (tensorrt/pycuda/numpy) + +cd tensorrtx/lenet + +python lenet.py -s # serialize model to plan file, i.e. 'lenet5.engine' + +python lenet.py -d # deserialize plan file and run inference + +# 4. see if the output is same as pytorchx/lenet +``` diff --git a/lenet/lenet.py b/lenet/lenet.py new file mode 100644 index 0000000..593f7bf --- /dev/null +++ b/lenet/lenet.py @@ -0,0 +1,190 @@ +import argparse +import os +import struct +import sys + +import numpy as np +import pycuda.autoinit +import pycuda.driver as cuda +import tensorrt as trt + +INPUT_H = 32 +INPUT_W = 32 +OUTPUT_SIZE = 10 +INPUT_BLOB_NAME = "data" +OUTPUT_BLOB_NAME = "prob" + +weight_path = "./lenet5.wts" +engine_path = "./lenet5.engine" + +gLogger = trt.Logger(trt.Logger.INFO) + + +def load_weights(file): + print(f"Loading weights: {file}") + + assert os.path.exists(file), 'Unable to load weight file.' + + weight_map = {} + with open(file, "r") as f: + lines = f.readlines() + count = int(lines[0]) + assert count == len(lines) - 1 + for i in range(1, count + 1): + splits = lines[i].split(" ") + name = splits[0] + cur_count = int(splits[1]) + assert cur_count + 2 == len(splits) + values = [] + for j in range(2, len(splits)): + # hex string to bytes to float + values.append(struct.unpack(">f", bytes.fromhex(splits[j]))) + weight_map[name] = np.array(values, dtype=np.float32) + + return weight_map + + +def createLenetEngine(maxBatchSize, builder, config, dt): + weight_map = load_weights(weight_path) + network = builder.create_network() + + data = network.add_input(INPUT_BLOB_NAME, dt, (1, INPUT_H, INPUT_W)) + assert data + + conv1 = network.add_convolution(input=data, + num_output_maps=6, + kernel_shape=(5, 5), + kernel=weight_map["conv1.weight"], + bias=weight_map["conv1.bias"]) + assert conv1 + conv1.stride = (1, 1) + + relu1 = network.add_activation(conv1.get_output(0), + type=trt.ActivationType.RELU) + assert relu1 + + pool1 = network.add_pooling(input=relu1.get_output(0), + window_size=trt.DimsHW(2, 2), + type=trt.PoolingType.AVERAGE) + assert pool1 + pool1.stride = (2, 2) + + conv2 = network.add_convolution(pool1.get_output(0), 16, trt.DimsHW(5, 5), + weight_map["conv2.weight"], + weight_map["conv2.bias"]) + assert conv2 + conv2.stride = (1, 1) + + relu2 = network.add_activation(conv2.get_output(0), + type=trt.ActivationType.RELU) + assert relu2 + + pool2 = network.add_pooling(input=relu2.get_output(0), + window_size=trt.DimsHW(2, 2), + type=trt.PoolingType.AVERAGE) + assert pool2 + pool2.stride = (2, 2) + + fc1 = network.add_fully_connected(input=pool2.get_output(0), + num_outputs=120, + kernel=weight_map['fc1.weight'], + bias=weight_map['fc1.bias']) + assert fc1 + + relu3 = network.add_activation(fc1.get_output(0), + type=trt.ActivationType.RELU) + assert relu3 + + fc2 = network.add_fully_connected(input=relu3.get_output(0), + num_outputs=84, + kernel=weight_map['fc2.weight'], + bias=weight_map['fc2.bias']) + assert fc2 + + relu4 = network.add_activation(fc2.get_output(0), + type=trt.ActivationType.RELU) + assert relu4 + + fc3 = network.add_fully_connected(input=relu4.get_output(0), + num_outputs=OUTPUT_SIZE, + kernel=weight_map['fc3.weight'], + bias=weight_map['fc3.bias']) + assert fc3 + + prob = network.add_softmax(fc3.get_output(0)) + assert prob + + prob.get_output(0).name = OUTPUT_BLOB_NAME + network.mark_output(prob.get_output(0)) + + # Build engine + builder.max_batch_size = maxBatchSize + builder.max_workspace_size = 1 << 20 + engine = builder.build_engine(network, config) + + del network + del weight_map + + return engine + + +def APIToModel(maxBatchSize): + builder = trt.Builder(gLogger) + config = builder.create_builder_config() + engine = createLenetEngine(maxBatchSize, builder, config, trt.float32) + assert engine + with open(engine_path, "wb") as f: + f.write(engine.serialize()) + + del engine + del builder + + +def doInference(context, host_in, host_out, batchSize): + engine = context.engine + assert engine.num_bindings == 2 + + devide_in = cuda.mem_alloc(host_in.nbytes) + devide_out = cuda.mem_alloc(host_out.nbytes) + bindings = [int(devide_in), int(devide_out)] + stream = cuda.Stream() + + cuda.memcpy_htod_async(devide_in, host_in, stream) + context.execute_async(bindings=bindings, stream_handle=stream.handle) + cuda.memcpy_dtoh_async(host_out, devide_out, stream) + stream.synchronize() + + +if __name__ == '__main__': + parser = argparse.ArgumentParser() + parser.add_argument("-s", action='store_true') + parser.add_argument("-d", action='store_true') + args = parser.parse_args() + + if not (args.s ^ args.d): + print("arguments not right!") + print("python lenet.py -s # serialize model to plan file") + print("python lenet.py -d # deserialize plan file and run inference") + sys.exit() + + if args.s: + APIToModel(1) + else: + runtime = trt.Runtime(gLogger) + assert runtime + + with open(engine_path, "rb") as f: + engine = runtime.deserialize_cuda_engine(f.read()) + assert engine + + context = engine.create_execution_context() + assert context + + data = np.ones((INPUT_H * INPUT_W), dtype=np.float32) + host_in = cuda.pagelocked_empty(INPUT_H * INPUT_W, dtype=np.float32) + np.copyto(host_in, data.ravel()) + host_out = cuda.pagelocked_empty(OUTPUT_SIZE, dtype=np.float32) + + doInference(context, host_in, host_out, 1) + + print(f'Output: {host_out}')