From 7f17ce4f9b6b6dce345d66736d38ccea576bdb27 Mon Sep 17 00:00:00 2001 From: irvingzhang0512 Date: Thu, 15 Apr 2021 13:17:41 +0800 Subject: [PATCH] Build ResNet50 by TensorRT Python Network Definition API (#482) * add tensorrt python api sample for resnet50 * update docs * typo * fix typo Co-authored-by: Wang Xinyu --- resnet/README.md | 22 +++- resnet/resnet50.py | 295 +++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 314 insertions(+), 3 deletions(-) create mode 100644 resnet/resnet50.py diff --git a/resnet/README.md b/resnet/README.md index 900fa66..d6fa246 100644 --- a/resnet/README.md +++ b/resnet/README.md @@ -1,13 +1,14 @@ # resnet -ResNet-18 adn ResNet-50 model from -"Deep Residual Learning for Image Recognition" +ResNet-18 and ResNet-50 model from "Deep Residual Learning for Image Recognition" For the Pytorch implementation, you can refer to [pytorchx/resnet](https://github.com/wang-xinyu/pytorchx/tree/master/resnet) Following tricks are used in this resnet, nothing special, residual connection and batchnorm are used. -- Batchnorm layer, implemented by scale layer. +- Batchnorm layer, implemented with scale layer. + +## TensorRT C++ API ``` // 1. generate resnet18.wts or resnet50.wts from [pytorchx/resnet](https://github.com/wang-xinyu/pytorchx/tree/master/resnet) @@ -38,4 +39,19 @@ sudo ./resnet50 -d // deserialize plan file and run inference // 4. see if the output is same as pytorchx/resnet ``` +### TensorRT Python API +``` +# 1. generate resnet50.wts from [pytorchx/resnet](https://github.com/wang-xinyu/pytorchx/tree/master/resnet) + +# 2. put resnet50.wts into tensorrtx/resnet + +# 3. install Python dependencies (tensorrt/pycuda/numpy) + +cd tensorrtx/resnet + +python resnet50.py -s // serialize model to plan file i.e. 'resnet50.engine' +python resnet50.py -d // deserialize plan file and run inference + +# 4. see if the output is same as pytorchx/resnet +``` diff --git a/resnet/resnet50.py b/resnet/resnet50.py new file mode 100644 index 0000000..fa3c84a --- /dev/null +++ b/resnet/resnet50.py @@ -0,0 +1,295 @@ +import argparse +import os +import struct +import sys + +import numpy as np +import pycuda.autoinit # noqa +import pycuda.driver as cuda +import tensorrt as trt + +BATCH_SIZE = 1 +INPUT_H = 224 +INPUT_W = 224 +OUTPUT_SIZE = 1000 +INPUT_BLOB_NAME = "data" +OUTPUT_BLOB_NAME = "prob" +EPS = 1e-5 + +WEIGHT_PATH = "./resnet50.wts" +ENGINE_PATH = "./resnet50.engine" + +TRT_LOGGER = 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 addBatchNorm2d(network, weight_map, input, layer_name, eps): + gamma = weight_map[layer_name + ".weight"] + beta = weight_map[layer_name + ".bias"] + mean = weight_map[layer_name + ".running_mean"] + var = weight_map[layer_name + ".running_var"] + var = np.sqrt(var + eps) + + scale = gamma / var + shift = -mean / var * gamma + beta + return network.add_scale(input=input, + mode=trt.ScaleMode.CHANNEL, + shift=shift, + scale=scale) + + +def bottleneck(network, weight_map, input, in_channels, out_channels, stride, + layer_name): + + conv1 = network.add_convolution(input=input, + num_output_maps=out_channels, + kernel_shape=(1, 1), + kernel=weight_map[layer_name + + "conv1.weight"], + bias=trt.Weights()) + assert conv1 + + bn1 = addBatchNorm2d(network, weight_map, conv1.get_output(0), + layer_name + "bn1", EPS) + assert bn1 + + relu1 = network.add_activation(bn1.get_output(0), + type=trt.ActivationType.RELU) + assert relu1 + + conv2 = network.add_convolution(input=relu1.get_output(0), + num_output_maps=out_channels, + kernel_shape=(3, 3), + kernel=weight_map[layer_name + + "conv2.weight"], + bias=trt.Weights()) + assert conv2 + conv2.stride = (stride, stride) + conv2.padding = (1, 1) + + bn2 = addBatchNorm2d(network, weight_map, conv2.get_output(0), + layer_name + "bn2", EPS) + assert bn2 + + relu2 = network.add_activation(bn2.get_output(0), + type=trt.ActivationType.RELU) + assert relu2 + + conv3 = network.add_convolution(input=relu2.get_output(0), + num_output_maps=out_channels * 4, + kernel_shape=(1, 1), + kernel=weight_map[layer_name + + "conv3.weight"], + bias=trt.Weights()) + assert conv3 + + bn3 = addBatchNorm2d(network, weight_map, conv3.get_output(0), + layer_name + "bn3", EPS) + assert bn3 + + if stride != 1 or in_channels != 4 * out_channels: + conv4 = network.add_convolution( + input=input, + num_output_maps=out_channels * 4, + kernel_shape=(1, 1), + kernel=weight_map[layer_name + "downsample.0.weight"], + bias=trt.Weights()) + assert conv4 + conv4.stride = (stride, stride) + + bn4 = addBatchNorm2d(network, weight_map, conv4.get_output(0), + layer_name + "downsample.1", EPS) + assert bn4 + + ew1 = network.add_elementwise(bn4.get_output(0), bn3.get_output(0), + trt.ElementWiseOperation.SUM) + else: + ew1 = network.add_elementwise(input, bn3.get_output(0), + trt.ElementWiseOperation.SUM) + assert ew1 + + relu3 = network.add_activation(ew1.get_output(0), + type=trt.ActivationType.RELU) + assert relu3 + + return relu3 + + +def createLenetEngine(maxBatchSize, builder, config, dt): + weight_map = load_weights(WEIGHT_PATH) + network = builder.create_network() + + data = network.add_input(INPUT_BLOB_NAME, dt, (3, INPUT_H, INPUT_W)) + assert data + + conv1 = network.add_convolution(input=data, + num_output_maps=64, + kernel_shape=(7, 7), + kernel=weight_map["conv1.weight"], + bias=trt.Weights()) + assert conv1 + conv1.stride = (2, 2) + conv1.padding = (3, 3) + + bn1 = addBatchNorm2d(network, weight_map, conv1.get_output(0), "bn1", EPS) + assert bn1 + + relu1 = network.add_activation(bn1.get_output(0), + type=trt.ActivationType.RELU) + assert relu1 + + pool1 = network.add_pooling(input=relu1.get_output(0), + window_size=trt.DimsHW(3, 3), + type=trt.PoolingType.MAX) + assert pool1 + pool1.stride = (2, 2) + pool1.padding = (1, 1) + + x = bottleneck(network, weight_map, pool1.get_output(0), 64, 64, 1, + "layer1.0.") + x = bottleneck(network, weight_map, x.get_output(0), 256, 64, 1, + "layer1.1.") + x = bottleneck(network, weight_map, x.get_output(0), 256, 64, 1, + "layer1.2.") + + x = bottleneck(network, weight_map, x.get_output(0), 256, 128, 2, + "layer2.0.") + x = bottleneck(network, weight_map, x.get_output(0), 512, 128, 1, + "layer2.1.") + x = bottleneck(network, weight_map, x.get_output(0), 512, 128, 1, + "layer2.2.") + x = bottleneck(network, weight_map, x.get_output(0), 512, 128, 1, + "layer2.3.") + + x = bottleneck(network, weight_map, x.get_output(0), 512, 256, 2, + "layer3.0.") + x = bottleneck(network, weight_map, x.get_output(0), 1024, 256, 1, + "layer3.1.") + x = bottleneck(network, weight_map, x.get_output(0), 1024, 256, 1, + "layer3.2.") + x = bottleneck(network, weight_map, x.get_output(0), 1024, 256, 1, + "layer3.3.") + x = bottleneck(network, weight_map, x.get_output(0), 1024, 256, 1, + "layer3.4.") + x = bottleneck(network, weight_map, x.get_output(0), 1024, 256, 1, + "layer3.5.") + + x = bottleneck(network, weight_map, x.get_output(0), 1024, 512, 2, + "layer4.0.") + x = bottleneck(network, weight_map, x.get_output(0), 2048, 512, 1, + "layer4.1.") + x = bottleneck(network, weight_map, x.get_output(0), 2048, 512, 1, + "layer4.2.") + + pool2 = network.add_pooling(x.get_output(0), + window_size=trt.DimsHW(7, 7), + type=trt.PoolingType.AVERAGE) + assert pool2 + pool2.stride = (1, 1) + + fc1 = network.add_fully_connected(input=pool2.get_output(0), + num_outputs=OUTPUT_SIZE, + kernel=weight_map['fc.weight'], + bias=weight_map['fc.bias']) + assert fc1 + + fc1.get_output(0).name = OUTPUT_BLOB_NAME + network.mark_output(fc1.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(TRT_LOGGER) + 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!\n" + "python resnet50.py -s # serialize model to plan file\n" + "python resnet50.py -d # deserialize plan file and run inference" + ) + sys.exit() + + if args.s: + APIToModel(BATCH_SIZE) + else: + runtime = trt.Runtime(TRT_LOGGER) + 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((BATCH_SIZE * 3 * INPUT_H * INPUT_W), dtype=np.float32) + host_in = cuda.pagelocked_empty(BATCH_SIZE * 3 * 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, BATCH_SIZE) + + print(f'Output: \n{host_out[:10]}\n{host_out[-10:]}')