From d9e2cbc294ae61bcd6b877156756cb06c057e9b1 Mon Sep 17 00:00:00 2001 From: SsisyphusTao <47047345+SsisyphusTao@users.noreply.github.com> Date: Sun, 23 May 2021 16:29:14 +0800 Subject: [PATCH] add centernet dla34 ctdet task. (#558) * add centernet dla34 ctdet task. * update readme. * update readme and fix a bug in sample. Co-authored-by: chandler --- centernet/README.md | 29 ++ centernet/centernet.py | 333 ++++++++++++++++ centernet/dcnv2Plugin/CMakeLists.txt | 21 + centernet/dcnv2Plugin/dcn_v2_im2col_cuda.cu | 399 +++++++++++++++++++ centernet/dcnv2Plugin/dcn_v2_im2col_cuda.h | 100 +++++ centernet/dcnv2Plugin/dcnv2Plugin.cpp | 402 ++++++++++++++++++++ centernet/dcnv2Plugin/dcnv2Plugin.h | 152 ++++++++ centernet/sample/common.py | 238 ++++++++++++ centernet/sample/test.py | 134 +++++++ 9 files changed, 1808 insertions(+) create mode 100644 centernet/README.md create mode 100644 centernet/centernet.py create mode 100644 centernet/dcnv2Plugin/CMakeLists.txt create mode 100644 centernet/dcnv2Plugin/dcn_v2_im2col_cuda.cu create mode 100644 centernet/dcnv2Plugin/dcn_v2_im2col_cuda.h create mode 100644 centernet/dcnv2Plugin/dcnv2Plugin.cpp create mode 100644 centernet/dcnv2Plugin/dcnv2Plugin.h create mode 100644 centernet/sample/common.py create mode 100644 centernet/sample/test.py diff --git a/centernet/README.md b/centernet/README.md new file mode 100644 index 0000000..b3661d7 --- /dev/null +++ b/centernet/README.md @@ -0,0 +1,29 @@ +# CenterNet + +This is the trt implementation of detection model [ctdet_coco_dla_2x](https://drive.google.com/open?id=1pl_-ael8wERdUREEnaIfqOV_VF2bEVRT) from [xingyizhou/CenterNet](https://github.com/xingyizhou/CenterNet) official work. + +## How to Run + +1. Follow [NVIDIA/TensorRT](https://github.com/NVIDIA/TensorRT) tutorial to build TensorRT7 + +2. Copy folder `dcnv2Plugin` to `TensorRT/plugin` and edit `InferPlugin.cpp` and `CMakeLists.txt` + +3. Rebuild to install custom plugin + +4. Use `tensorrt-7.2.3.4-cp36-none-linux_x86_64.whl` in TensorRT OSS to update your python-tensorrt + +5. Run `python centernet.py -m ${PTH_PATH} -s` to create trt engine + +## Sample + +``` +// Download ctdet_coco_dla_2x.pth and transfer it into trt engine first +// Download the test img from https://raw.githubusercontent.com/tensorflow/models/master/research/deeplab/g3doc/img/image2.jpg or choose your own one +cd sample +python test.py ${ENGINE_PATH} ${IMG_PATH} +``` +![trt_out](https://user-images.githubusercontent.com/47047345/119128637-7a878900-ba68-11eb-91ff-5dcc10f01b77.jpg) + +## TODO + +Integrate the post process with trt engine to make it more easier to use. \ No newline at end of file diff --git a/centernet/centernet.py b/centernet/centernet.py new file mode 100644 index 0000000..197ee85 --- /dev/null +++ b/centernet/centernet.py @@ -0,0 +1,333 @@ +import numpy as np + +import tensorrt as trt +import torch + +from sample import common +import argparse +import time + +# You can set the logger severity higher to suppress messages (or lower to display more messages). +TRT_LOGGER = trt.Logger(trt.Logger.WARNING) + +trt.init_libnvinfer_plugins(TRT_LOGGER, '') +PLUGIN_CREATORS = trt.get_plugin_registry().plugin_creator_list + +for plugin_creator in PLUGIN_CREATORS: + if plugin_creator.name == 'DCNv2_TRT': + dcnCreator = plugin_creator + + +class ModelData(object): + INPUT_NAME = "data" + INPUT_SHAPE = (3, 512, 512) + OUTPUT_NAME = "prob" + DTYPE = trt.float16 + + +class Centernet_dla34(object): + def __init__(self, weights) -> None: + super().__init__() + self.weights = weights + self.levels = [1, 1, 1, 2, 2, 1] + self.channels = [16, 32, 64, 128, 256, 512] + self.down_ratio = 4 + self.last_level = 5 + self.engine = self.build_engine() + + def add_batchnorm_2d(self, input_tensor, parent): + gamma = self.weights[parent + '.weight'].numpy() + beta = self.weights[parent + '.bias'].numpy() + mean = self.weights[parent + '.running_mean'].numpy() + var = self.weights[parent + '.running_var'].numpy() + eps = 1e-5 + + scale = gamma / np.sqrt(var + eps) + shift = beta - mean * gamma / np.sqrt(var + eps) + power = np.ones_like(scale) + + return self.network.add_scale(input=input_tensor.get_output(0), mode=trt.ScaleMode.CHANNEL, shift=shift, scale=scale, power=power) + + def add_basic_block(self, input_tensor, out_channels, residual=None, stride=1, dilation=1, parent=''): + conv1_w = self.weights[parent + '.conv1.weight'].numpy() + conv1 = self.network.add_convolution(input=input_tensor.get_output( + 0), num_output_maps=out_channels, kernel_shape=(3, 3), kernel=conv1_w) + conv1.stride = (stride, stride) + conv1.padding = (dilation, dilation) + conv1.dilation = (dilation, dilation) + + bn1 = self.add_batchnorm_2d(conv1, parent + '.bn1') + ac1 = self.network.add_activation( + input=bn1.get_output(0), type=trt.ActivationType.RELU) + + conv2_w = self.weights[parent + '.conv2.weight'].numpy() + conv2 = self.network.add_convolution(input=ac1.get_output( + 0), num_output_maps=out_channels, kernel_shape=(3, 3), kernel=conv2_w) + conv2.padding = (dilation, dilation) + conv2.dilation = (dilation, dilation) + + out = self.add_batchnorm_2d(conv2, parent + '.bn2') + + if residual is None: + out = self.network.add_elementwise(input_tensor.get_output( + 0), out.get_output(0), trt.ElementWiseOperation.SUM) + else: + out = self.network.add_elementwise(residual.get_output( + 0), out.get_output(0), trt.ElementWiseOperation.SUM) + return self.network.add_activation(input=out.get_output(0), type=trt.ActivationType.RELU) + + def add_level(self, input_tensor, out_channels, stride=1, dilation=1, parent=''): + conv1_w = self.weights[parent + '.0.weight'].numpy() + conv1 = self.network.add_convolution(input=input_tensor.get_output( + 0), num_output_maps=out_channels, kernel_shape=(3, 3), kernel=conv1_w) + conv1.stride = (stride, stride) + conv1.padding = (dilation, dilation) + conv1.dilation = (dilation, dilation) + + bn1 = self.add_batchnorm_2d(conv1, parent + '.1') + ac1 = self.network.add_activation( + input=bn1.get_output(0), type=trt.ActivationType.RELU) + return ac1 + + def add_root(self, input_tensors: list, out_channels, kernel_size=1, residual=False, parent=''): + ct = self.network.add_concatenation( + [x.get_output(0) for x in input_tensors]) + + conv_w = self.weights[parent + '.conv.weight'].numpy() + conv = self.network.add_convolution(input=ct.get_output( + 0), num_output_maps=out_channels, kernel_shape=(1, 1), kernel=conv_w) + conv.padding = ((kernel_size - 1) // 2, (kernel_size - 1) // 2) + + bn1 = self.add_batchnorm_2d(conv, parent + '.bn') + out = self.network.add_activation( + input=bn1.get_output(0), type=trt.ActivationType.RELU) + + if residual: + out = self.network.add_elementwise(input_tensors[0].get_output( + 0), out.get_output(0), trt.ElementWiseOperation.SUM) + + return self.network.add_activation(input=out.get_output(0), type=trt.ActivationType.RELU) + + def add_tree(self, input_tensor, level, out_channels, residual=None, children=None, stride=1, level_root=False, parent=''): + children = [] if children is None else children + if stride > 1: + bottom = self.network.add_pooling(input_tensor.get_output( + 0), trt.PoolingType.MAX, (stride, stride)) + bottom.stride = (stride, stride) + else: + bottom = input_tensor + + if input_tensor.get_output(0).shape[0] != out_channels: + project_conv1_w = self.weights[parent + + '.project.0.weight'].numpy() + project_conv1 = self.network.add_convolution(input=bottom.get_output( + 0), num_output_maps=out_channels, kernel_shape=(1, 1), kernel=project_conv1_w) + residual = self.add_batchnorm_2d( + project_conv1, parent + '.project.1') + else: + residual = bottom + + if level_root: + children.append(bottom) + + if level == 1: + tree1 = self.add_basic_block( + input_tensor, out_channels, residual, stride, parent=parent+'.tree1') + tree2 = self.add_basic_block( + tree1, out_channels, parent=parent+'.tree2') + return self.add_root([tree2, tree1]+children, out_channels, parent=parent+'.root') + else: + tree1 = self.add_tree(input_tensor, level-1, out_channels, + residual, stride=stride, parent=parent+'.tree1') + children.append(tree1) + return self.add_tree(tree1, level-1, out_channels, children=children, parent=parent+'.tree2') + + def add_base(self, input_tensor, parent): + base_conv1_w = self.weights[parent+'.base_layer.0.weight'].numpy() + base_conv1 = self.network.add_convolution( + input=input_tensor, num_output_maps=self.channels[0], kernel_shape=(7, 7), kernel=base_conv1_w) + base_conv1.padding = (3, 3) + + base_bn1 = self.add_batchnorm_2d(base_conv1, parent+'.base_layer.1') + base_ac1 = self.network.add_activation( + input=base_bn1.get_output(0), type=trt.ActivationType.RELU) + + level0 = self.add_level( + base_ac1, self.channels[0], parent=parent+'.level0') + level1 = self.add_level( + level0, self.channels[1], 2, parent=parent+'.level1') + + level2 = self.add_tree( + level1, self.levels[2], self.channels[2], stride=2, level_root=False, parent=parent+'.level2') + level3 = self.add_tree( + level2, self.levels[3], self.channels[3], stride=2, level_root=True, parent=parent+'.level3') + level4 = self.add_tree( + level3, self.levels[4], self.channels[4], stride=2, level_root=True, parent=parent+'.level4') + level5 = self.add_tree( + level4, self.levels[5], self.channels[5], stride=2, level_root=True, parent=parent+'.level5') + + return [level0, level1, level2, level3, level4, level5] + + def add_deform_conv(self, input_tensor, out_channels, kernel=3, stride=1, padding=1, dilation=1, deformable_group=1, parent=''): + conv_offset_mask_w = self.weights[parent + + '.conv.conv_offset_mask.weight'].numpy() + conv_offset_mask_b = self.weights[parent + + '.conv.conv_offset_mask.bias'].numpy() + conv_offset_mask = self.network.add_convolution(input=input_tensor.get_output(0), + num_output_maps=deformable_group*3*kernel*kernel, + kernel_shape=( + kernel, kernel), + kernel=conv_offset_mask_w, + bias=conv_offset_mask_b) + conv_offset_mask.stride = (stride, stride) + conv_offset_mask.padding = (padding, padding) + + out_channels = trt.PluginField("out_channels", np.array( + [out_channels], dtype=np.int32), trt.PluginFieldType.INT32) + kernel = trt.PluginField("kernel", np.array( + [kernel], dtype=np.int32), trt.PluginFieldType.INT32) + deformable_group = trt.PluginField("deformable_group", np.array( + [deformable_group], dtype=np.int32), trt.PluginFieldType.INT32) + dilation = trt.PluginField("dilation", np.array( + [dilation], dtype=np.int32), trt.PluginFieldType.INT32) + padding = trt.PluginField("padding", np.array( + [padding], dtype=np.int32), trt.PluginFieldType.INT32) + stride = trt.PluginField("stride", np.array( + [stride], dtype=np.int32), trt.PluginFieldType.INT32) + weight = trt.PluginField( + "weight", self.weights[parent + '.conv.weight'].numpy(), trt.PluginFieldType.FLOAT32) + bias = trt.PluginField( + "bias", self.weights[parent + '.conv.bias'].numpy(), trt.PluginFieldType.FLOAT32) + field_collection = trt.PluginFieldCollection( + [out_channels, kernel, deformable_group, dilation, padding, stride, weight, bias]) + DCN = dcnCreator.create_plugin( + name='DCNv2_TRT', field_collection=field_collection) + + sigmoid_conv_offset_mask = self.network.add_activation( + input=conv_offset_mask.get_output(0), type=trt.ActivationType.SIGMOID) + + dcn = self.network.add_plugin_v2(inputs=[input_tensor.get_output( + 0), conv_offset_mask.get_output(0), sigmoid_conv_offset_mask.get_output(0)], plugin=DCN) + bn = self.add_batchnorm_2d(dcn, parent+'.actf.0') + return self.network.add_activation(input=bn.get_output(0), type=trt.ActivationType.RELU) + + def add_ida_up(self, input_tensors, out_channels, up_f, startp, parent): + for i in range(startp + 1, len(input_tensors)): + proj = self.add_deform_conv( + input_tensors[i], out_channels, parent=parent+'.proj_%d' % (i-startp)) + f = up_f[i-startp] + up_w = self.weights[parent + '.up_%d.weight' % (i-startp)].numpy() + up = self.network.add_deconvolution( + proj.get_output(0), out_channels, (f*2, f*2), up_w) + up.stride = (f, f) + up.padding = (f//2, f//2) + up.num_groups = out_channels + node = self.network.add_elementwise( + input_tensors[i-1].get_output(0), up.get_output(0), trt.ElementWiseOperation.SUM) + input_tensors[i] = self.add_deform_conv( + node, out_channels, parent=parent+'.node_%d' % (i-startp)) + return input_tensors + + def add_dla_up(self, input_tensors, first_level, parent): + channels = self.channels[first_level:] + scales = [2 ** i for i in range(len(self.channels[first_level:]))] + scales = np.array(scales, dtype=int) + out = [input_tensors[-1]] + for i in range(len(channels) - 1): + j = -i - 2 + input_tensors = self.add_ida_up( + input_tensors, channels[j], scales[j:] // scales[j], len(input_tensors) - i - 2, parent+'.ida_%d' % i) + out.insert(0, input_tensors[-1]) + scales[j + 1:] = scales[j] + channels[j + 1:] = [channels[j] for _ in channels[j + 1:]] + return out + + def add_head(self, input_tensor, out_channels, head, head_conv=256, final_kernal=1): + conv1_w = self.weights[head+'.0.weight'].numpy() + conv1_b = self.weights[head+'.0.bias'].numpy() + conv1 = self.network.add_convolution( + input_tensor.get_output(0), head_conv, (3, 3), conv1_w, conv1_b) + conv1.padding = (1, 1) + ac1 = self.network.add_activation( + input=conv1.get_output(0), type=trt.ActivationType.RELU) + conv2_w = self.weights[head + '.2.weight'].numpy() + conv2_b = self.weights[head+'.2.bias'].numpy() + conv2 = self.network.add_convolution(ac1.get_output( + 0), out_channels, (final_kernal, final_kernal), conv2_w, conv2_b) + return conv2 + + def populate_network(self): + # Configure the network layers based on the self.weights provided. + input_tensor = self.network.add_input( + name=ModelData.INPUT_NAME, dtype=ModelData.DTYPE, shape=ModelData.INPUT_SHAPE) + + y = self.add_base(input_tensor, 'module.base') + + first_level = int(np.log2(self.down_ratio)) + last_level = self.last_level + dla_up = self.add_dla_up(y, first_level, 'module.dla_up') + ida_up = self.add_ida_up(dla_up[:last_level-first_level], self.channels[first_level], [ + 2 ** i for i in range(last_level - first_level)], 0, 'module.ida_up') + + hm = self.add_head(ida_up[-1], 80, 'module.hm') + wh = self.add_head(ida_up[-1], 2, 'module.wh') + reg = self.add_head(ida_up[-1], 2, 'module.reg') + + hm.get_output(0).name = 'hm' + wh.get_output(0).name = 'wh' + reg.get_output(0).name = 'reg' + self.network.mark_output(tensor=hm.get_output(0)) + self.network.mark_output(tensor=wh.get_output(0)) + self.network.mark_output(tensor=reg.get_output(0)) + + def build_engine(self): + # For more information on TRT basics, refer to the introductory samples. + with trt.Builder(TRT_LOGGER) as builder, builder.create_network() as network: + self.network = network + builder.max_workspace_size = common.GiB(1) + builder.max_batch_size = 1 + # Populate the network using self.weights from the PyTorch model. + self.populate_network() + # Build and return an engine. + return builder.build_cuda_engine(self.network) + + +def load_random_test_case(pagelocked_buffer): + # Select an image at random to be the test case. + img = np.random.randn(1, 3, 512, 512).astype(np.float32) + # Copy to the pagelocked input buffer + np.copyto(pagelocked_buffer, img.ravel()) + return img + + +def main(args): + # Get the PyTorch weights + weights = torch.load(args.model, map_location={ + 'cuda:0': 'cpu'})['state_dict'] + # Do inference with TensorRT. + with Centernet_dla34(weights).engine as engine: + if args.save_engine: + with open('centernet.engine', "wb") as f: + f.write(engine.serialize()) + inputs, outputs, bindings, stream = common.allocate_buffers(engine) + with engine.create_execution_context() as context: + img = load_random_test_case(pagelocked_buffer=inputs[0].host) + # For more information on performing inference, refer to the introductory samples. + # The common.do_inference function will return a list of outputs - we only have one in this case. + t = time.time() + [hm, wh, reg] = common.do_inference( + context, bindings=bindings, inputs=inputs, outputs=outputs, stream=stream, batch_size=1) + t = time.time() - t + print('output: hm:%f, wh:%f, reg:%f' % + (hm.mean(), wh.mean(), reg.mean())) + print(t) + + +if __name__ == '__main__': + parser = argparse.ArgumentParser(description='CenterNet dla34 ctdet') + parser.add_argument('--model', '-m', type=str, + default='./ctdet_coco_dla_2x.pth', help='path of pytorch .pth') + parser.add_argument('--save_engine', '-s', + action='store_true', help='if save trt engine') + args = parser.parse_args() + main(args) diff --git a/centernet/dcnv2Plugin/CMakeLists.txt b/centernet/dcnv2Plugin/CMakeLists.txt new file mode 100644 index 0000000..435d9a1 --- /dev/null +++ b/centernet/dcnv2Plugin/CMakeLists.txt @@ -0,0 +1,21 @@ +# +# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +file(GLOB SRCS *.cpp) +set(PLUGIN_SOURCES ${PLUGIN_SOURCES} ${SRCS}) +set(PLUGIN_SOURCES ${PLUGIN_SOURCES} PARENT_SCOPE) +file(GLOB CU_SRCS *.cu) +set(PLUGIN_CU_SOURCES ${PLUGIN_CU_SOURCES} ${CU_SRCS}) +set(PLUGIN_CU_SOURCES ${PLUGIN_CU_SOURCES} PARENT_SCOPE) \ No newline at end of file diff --git a/centernet/dcnv2Plugin/dcn_v2_im2col_cuda.cu b/centernet/dcnv2Plugin/dcn_v2_im2col_cuda.cu new file mode 100644 index 0000000..43271bc --- /dev/null +++ b/centernet/dcnv2Plugin/dcn_v2_im2col_cuda.cu @@ -0,0 +1,399 @@ +#include "dcn_v2_im2col_cuda.h" +#include +#include +#include + +#define CUDA_KERNEL_LOOP(i, n) \ + for (int i = blockIdx.x * blockDim.x + threadIdx.x; \ + i < (n); \ + i += blockDim.x * gridDim.x) + +const int CUDA_NUM_THREADS = 512; +//inline int GET_BLOCKS(const int N) +//{ +// return (N + CUDA_NUM_THREADS - 1) / CUDA_NUM_THREADS; +//} +dim3 GET_BLOCKS(uint n) +{ + uint k = (n - 1) /CUDA_NUM_THREADS + 1; + uint x = k ; + uint y = 1 ; + if (x > 65535 ) + { + x = ceil(sqrt(x)); + y = (n - 1 )/(x*CUDA_NUM_THREADS) + 1; + } + dim3 d = {x,y,1} ; + return d; +} + +__device__ float dmcn_im2col_bilinear(const float *bottom_data, const int data_width, + const int height, const int width, float h, float w) +{ + int h_low = floor(h); + int w_low = floor(w); + int h_high = h_low + 1; + int w_high = w_low + 1; + + float lh = h - h_low; + float lw = w - w_low; + float hh = 1 - lh, hw = 1 - lw; + + float v1 = 0; + if (h_low >= 0 && w_low >= 0) + v1 = bottom_data[h_low * data_width + w_low]; + float v2 = 0; + if (h_low >= 0 && w_high <= width - 1) + v2 = bottom_data[h_low * data_width + w_high]; + float v3 = 0; + if (h_high <= height - 1 && w_low >= 0) + v3 = bottom_data[h_high * data_width + w_low]; + float v4 = 0; + if (h_high <= height - 1 && w_high <= width - 1) + v4 = bottom_data[h_high * data_width + w_high]; + + float w1 = hh * hw, w2 = hh * lw, w3 = lh * hw, w4 = lh * lw; + + float val = (w1 * v1 + w2 * v2 + w3 * v3 + w4 * v4); + return val; +} + +__device__ float dmcn_get_gradient_weight(float argmax_h, float argmax_w, + const int h, const int w, const int height, const int width) +{ + if (argmax_h <= -1 || argmax_h >= height || argmax_w <= -1 || argmax_w >= width) + { + //empty + return 0; + } + + int argmax_h_low = floor(argmax_h); + int argmax_w_low = floor(argmax_w); + int argmax_h_high = argmax_h_low + 1; + int argmax_w_high = argmax_w_low + 1; + + float weight = 0; + if (h == argmax_h_low && w == argmax_w_low) + weight = (h + 1 - argmax_h) * (w + 1 - argmax_w); + if (h == argmax_h_low && w == argmax_w_high) + weight = (h + 1 - argmax_h) * (argmax_w + 1 - w); + if (h == argmax_h_high && w == argmax_w_low) + weight = (argmax_h + 1 - h) * (w + 1 - argmax_w); + if (h == argmax_h_high && w == argmax_w_high) + weight = (argmax_h + 1 - h) * (argmax_w + 1 - w); + return weight; +} + +__device__ float dmcn_get_coordinate_weight(float argmax_h, float argmax_w, + const int height, const int width, const float *im_data, + const int data_width, const int bp_dir) +{ + if (argmax_h <= -1 || argmax_h >= height || argmax_w <= -1 || argmax_w >= width) + { + //empty + return 0; + } + + int argmax_h_low = floor(argmax_h); + int argmax_w_low = floor(argmax_w); + int argmax_h_high = argmax_h_low + 1; + int argmax_w_high = argmax_w_low + 1; + + float weight = 0; + + if (bp_dir == 0) + { + if (argmax_h_low >= 0 && argmax_w_low >= 0) + weight += -1 * (argmax_w_low + 1 - argmax_w) * im_data[argmax_h_low * data_width + argmax_w_low]; + if (argmax_h_low >= 0 && argmax_w_high <= width - 1) + weight += -1 * (argmax_w - argmax_w_low) * im_data[argmax_h_low * data_width + argmax_w_high]; + if (argmax_h_high <= height - 1 && argmax_w_low >= 0) + weight += (argmax_w_low + 1 - argmax_w) * im_data[argmax_h_high * data_width + argmax_w_low]; + if (argmax_h_high <= height - 1 && argmax_w_high <= width - 1) + weight += (argmax_w - argmax_w_low) * im_data[argmax_h_high * data_width + argmax_w_high]; + } + else if (bp_dir == 1) + { + if (argmax_h_low >= 0 && argmax_w_low >= 0) + weight += -1 * (argmax_h_low + 1 - argmax_h) * im_data[argmax_h_low * data_width + argmax_w_low]; + if (argmax_h_low >= 0 && argmax_w_high <= width - 1) + weight += (argmax_h_low + 1 - argmax_h) * im_data[argmax_h_low * data_width + argmax_w_high]; + if (argmax_h_high <= height - 1 && argmax_w_low >= 0) + weight += -1 * (argmax_h - argmax_h_low) * im_data[argmax_h_high * data_width + argmax_w_low]; + if (argmax_h_high <= height - 1 && argmax_w_high <= width - 1) + weight += (argmax_h - argmax_h_low) * im_data[argmax_h_high * data_width + argmax_w_high]; + } + + return weight; +} + +__global__ void modulated_deformable_im2col_gpu_kernel(const int n, + const float *data_im, const float *data_offset, const float *data_mask, + const int height, const int width, const int kernel_h, const int kernel_w, + const int pad_h, const int pad_w, + const int stride_h, const int stride_w, + const int dilation_h, const int dilation_w, + const int channel_per_deformable_group, + const int batch_size, const int num_channels, const int deformable_group, + const int height_col, const int width_col, + float *data_col) +{ + CUDA_KERNEL_LOOP(index, n) + { + // index index of output matrix + const int w_col = index % width_col; + const int h_col = (index / width_col) % height_col; + const int b_col = (index / width_col / height_col) % batch_size; + const int c_im = (index / width_col / height_col) / batch_size; + const int c_col = c_im * kernel_h * kernel_w; + + // compute deformable group index + const int deformable_group_index = c_im / channel_per_deformable_group; + + const int h_in = h_col * stride_h - pad_h; + const int w_in = w_col * stride_w - pad_w; + + float *data_col_ptr = data_col + ((c_col * batch_size + b_col) * height_col + h_col) * width_col + w_col; + //const float* data_im_ptr = data_im + ((b_col * num_channels + c_im) * height + h_in) * width + w_in; + const float *data_im_ptr = data_im + (b_col * num_channels + c_im) * height * width; + const float *data_offset_ptr = data_offset + (b_col * deformable_group + deformable_group_index) * 2 * kernel_h * kernel_w * height_col * width_col; + + const float *data_mask_ptr = data_mask + (b_col * deformable_group + deformable_group_index) * kernel_h * kernel_w * height_col * width_col; + + for (int i = 0; i < kernel_h; ++i) + { + for (int j = 0; j < kernel_w; ++j) + { + const int data_offset_h_ptr = ((2 * (i * kernel_w + j)) * height_col + h_col) * width_col + w_col; + const int data_offset_w_ptr = ((2 * (i * kernel_w + j) + 1) * height_col + h_col) * width_col + w_col; + const int data_mask_hw_ptr = ((i * kernel_w + j) * height_col + h_col) * width_col + w_col; + const float offset_h = data_offset_ptr[data_offset_h_ptr]; + const float offset_w = data_offset_ptr[data_offset_w_ptr]; + const float mask = data_mask_ptr[data_mask_hw_ptr]; + float val = static_cast(0); + const float h_im = h_in + i * dilation_h + offset_h; + const float w_im = w_in + j * dilation_w + offset_w; + //if (h_im >= 0 && w_im >= 0 && h_im < height && w_im < width) { + if (h_im > -1 && w_im > -1 && h_im < height && w_im < width) + { + //const float map_h = i * dilation_h + offset_h; + //const float map_w = j * dilation_w + offset_w; + //const int cur_height = height - h_in; + //const int cur_width = width - w_in; + //val = dmcn_im2col_bilinear(data_im_ptr, width, cur_height, cur_width, map_h, map_w); + val = dmcn_im2col_bilinear(data_im_ptr, width, height, width, h_im, w_im); + } + *data_col_ptr = val * mask; + data_col_ptr += batch_size * height_col * width_col; + //data_col_ptr += height_col * width_col; + } + } + } +} + +__global__ void modulated_deformable_col2im_gpu_kernel(const int n, + const float *data_col, const float *data_offset, const float *data_mask, + const int channels, const int height, const int width, + const int kernel_h, const int kernel_w, + const int pad_h, const int pad_w, + const int stride_h, const int stride_w, + const int dilation_h, const int dilation_w, + const int channel_per_deformable_group, + const int batch_size, const int deformable_group, + const int height_col, const int width_col, + float *grad_im) +{ + CUDA_KERNEL_LOOP(index, n) + { + const int j = (index / width_col / height_col / batch_size) % kernel_w; + const int i = (index / width_col / height_col / batch_size / kernel_w) % kernel_h; + const int c = index / width_col / height_col / batch_size / kernel_w / kernel_h; + // compute the start and end of the output + + const int deformable_group_index = c / channel_per_deformable_group; + + int w_out = index % width_col; + int h_out = (index / width_col) % height_col; + int b = (index / width_col / height_col) % batch_size; + int w_in = w_out * stride_w - pad_w; + int h_in = h_out * stride_h - pad_h; + + const float *data_offset_ptr = data_offset + (b * deformable_group + deformable_group_index) * 2 * kernel_h * kernel_w * height_col * width_col; + const float *data_mask_ptr = data_mask + (b * deformable_group + deformable_group_index) * kernel_h * kernel_w * height_col * width_col; + const int data_offset_h_ptr = ((2 * (i * kernel_w + j)) * height_col + h_out) * width_col + w_out; + const int data_offset_w_ptr = ((2 * (i * kernel_w + j) + 1) * height_col + h_out) * width_col + w_out; + const int data_mask_hw_ptr = ((i * kernel_w + j) * height_col + h_out) * width_col + w_out; + const float offset_h = data_offset_ptr[data_offset_h_ptr]; + const float offset_w = data_offset_ptr[data_offset_w_ptr]; + const float mask = data_mask_ptr[data_mask_hw_ptr]; + const float cur_inv_h_data = h_in + i * dilation_h + offset_h; + const float cur_inv_w_data = w_in + j * dilation_w + offset_w; + + const float cur_top_grad = data_col[index] * mask; + const int cur_h = (int)cur_inv_h_data; + const int cur_w = (int)cur_inv_w_data; + for (int dy = -2; dy <= 2; dy++) + { + for (int dx = -2; dx <= 2; dx++) + { + if (cur_h + dy >= 0 && cur_h + dy < height && + cur_w + dx >= 0 && cur_w + dx < width && + abs(cur_inv_h_data - (cur_h + dy)) < 1 && + abs(cur_inv_w_data - (cur_w + dx)) < 1) + { + int cur_bottom_grad_pos = ((b * channels + c) * height + cur_h + dy) * width + cur_w + dx; + float weight = dmcn_get_gradient_weight(cur_inv_h_data, cur_inv_w_data, cur_h + dy, cur_w + dx, height, width); + atomicAdd(grad_im + cur_bottom_grad_pos, weight * cur_top_grad); + } + } + } + } +} + +__global__ void modulated_deformable_col2im_coord_gpu_kernel(const int n, + const float *data_col, const float *data_im, + const float *data_offset, const float *data_mask, + const int channels, const int height, const int width, + const int kernel_h, const int kernel_w, + const int pad_h, const int pad_w, + const int stride_h, const int stride_w, + const int dilation_h, const int dilation_w, + const int channel_per_deformable_group, + const int batch_size, const int offset_channels, const int deformable_group, + const int height_col, const int width_col, + float *grad_offset, float *grad_mask) +{ + CUDA_KERNEL_LOOP(index, n) + { + float val = 0, mval = 0; + int w = index % width_col; + int h = (index / width_col) % height_col; + int c = (index / width_col / height_col) % offset_channels; + int b = (index / width_col / height_col) / offset_channels; + // compute the start and end of the output + + const int deformable_group_index = c / (2 * kernel_h * kernel_w); + const int col_step = kernel_h * kernel_w; + int cnt = 0; + const float *data_col_ptr = data_col + deformable_group_index * channel_per_deformable_group * batch_size * width_col * height_col; + const float *data_im_ptr = data_im + (b * deformable_group + deformable_group_index) * channel_per_deformable_group / kernel_h / kernel_w * height * width; + const float *data_offset_ptr = data_offset + (b * deformable_group + deformable_group_index) * 2 * kernel_h * kernel_w * height_col * width_col; + const float *data_mask_ptr = data_mask + (b * deformable_group + deformable_group_index) * kernel_h * kernel_w * height_col * width_col; + + const int offset_c = c - deformable_group_index * 2 * kernel_h * kernel_w; + + for (int col_c = (offset_c / 2); col_c < channel_per_deformable_group; col_c += col_step) + { + const int col_pos = (((col_c * batch_size + b) * height_col) + h) * width_col + w; + const int bp_dir = offset_c % 2; + + int j = (col_pos / width_col / height_col / batch_size) % kernel_w; + int i = (col_pos / width_col / height_col / batch_size / kernel_w) % kernel_h; + int w_out = col_pos % width_col; + int h_out = (col_pos / width_col) % height_col; + int w_in = w_out * stride_w - pad_w; + int h_in = h_out * stride_h - pad_h; + const int data_offset_h_ptr = (((2 * (i * kernel_w + j)) * height_col + h_out) * width_col + w_out); + const int data_offset_w_ptr = (((2 * (i * kernel_w + j) + 1) * height_col + h_out) * width_col + w_out); + const int data_mask_hw_ptr = (((i * kernel_w + j) * height_col + h_out) * width_col + w_out); + const float offset_h = data_offset_ptr[data_offset_h_ptr]; + const float offset_w = data_offset_ptr[data_offset_w_ptr]; + const float mask = data_mask_ptr[data_mask_hw_ptr]; + float inv_h = h_in + i * dilation_h + offset_h; + float inv_w = w_in + j * dilation_w + offset_w; + if (inv_h <= -1 || inv_w <= -1 || inv_h >= height || inv_w >= width) + { + inv_h = inv_w = -2; + } + else + { + mval += data_col_ptr[col_pos] * dmcn_im2col_bilinear(data_im_ptr + cnt * height * width, width, height, width, inv_h, inv_w); + } + const float weight = dmcn_get_coordinate_weight( + inv_h, inv_w, + height, width, data_im_ptr + cnt * height * width, width, bp_dir); + val += weight * data_col_ptr[col_pos] * mask; + cnt += 1; + } + // KERNEL_ASSIGN(grad_offset[index], offset_req, val); + grad_offset[index] = val; + if (offset_c % 2 == 0) + // KERNEL_ASSIGN(grad_mask[(((b * deformable_group + deformable_group_index) * kernel_h * kernel_w + offset_c / 2) * height_col + h) * width_col + w], mask_req, mval); + grad_mask[(((b * deformable_group + deformable_group_index) * kernel_h * kernel_w + offset_c / 2) * height_col + h) * width_col + w] = mval; + } +} + +void modulated_deformable_im2col_cuda(cudaStream_t stream, + const float* data_im, const float* data_offset, const float* data_mask, + const int batch_size, const int channels, const int height_im, const int width_im, + const int height_col, const int width_col, const int kernel_h, const int kenerl_w, + const int pad_h, const int pad_w, const int stride_h, const int stride_w, + const int dilation_h, const int dilation_w, + const int deformable_group, float* data_col) { + // num_axes should be smaller than block size + const int channel_per_deformable_group = channels / deformable_group; + const int num_kernels = channels * batch_size * height_col * width_col; + modulated_deformable_im2col_gpu_kernel + <<>>( + num_kernels, data_im, data_offset, data_mask, height_im, width_im, kernel_h, kenerl_w, + pad_h, pad_w, stride_h, stride_w, dilation_h, dilation_w, channel_per_deformable_group, + batch_size, channels, deformable_group, height_col, width_col, data_col); + + cudaError_t err = cudaGetLastError(); + if (err != cudaSuccess) + { + printf("error in modulated_deformable_im2col_cuda: %s\n", cudaGetErrorString(err)); + } + +} + +void modulated_deformable_col2im_cuda(cudaStream_t stream, + const float* data_col, const float* data_offset, const float* data_mask, + const int batch_size, const int channels, const int height_im, const int width_im, + const int height_col, const int width_col, const int kernel_h, const int kernel_w, + const int pad_h, const int pad_w, const int stride_h, const int stride_w, + const int dilation_h, const int dilation_w, + const int deformable_group, float* grad_im){ + + const int channel_per_deformable_group = channels / deformable_group; + const int num_kernels = channels * kernel_h * kernel_w * batch_size * height_col * width_col; + modulated_deformable_col2im_gpu_kernel + <<>>( + num_kernels, data_col, data_offset, data_mask, channels, height_im, width_im, + kernel_h, kernel_w, pad_h, pad_h, stride_h, stride_w, + dilation_h, dilation_w, channel_per_deformable_group, + batch_size, deformable_group, height_col, width_col, grad_im); + cudaError_t err = cudaGetLastError(); + if (err != cudaSuccess) + { + printf("error in modulated_deformable_col2im_cuda: %s\n", cudaGetErrorString(err)); + } + +} + +void modulated_deformable_col2im_coord_cuda(cudaStream_t stream, + const float* data_col, const float* data_im, const float* data_offset, const float* data_mask, + const int batch_size, const int channels, const int height_im, const int width_im, + const int height_col, const int width_col, const int kernel_h, const int kernel_w, + const int pad_h, const int pad_w, const int stride_h, const int stride_w, + const int dilation_h, const int dilation_w, + const int deformable_group, + float* grad_offset, float* grad_mask) { + const int num_kernels = batch_size * height_col * width_col * 2 * kernel_h * kernel_w * deformable_group; + const int channel_per_deformable_group = channels * kernel_h * kernel_w / deformable_group; + modulated_deformable_col2im_coord_gpu_kernel + <<>>( + num_kernels, data_col, data_im, data_offset, data_mask, channels, height_im, width_im, + kernel_h, kernel_w, pad_h, pad_w, stride_h, stride_w, + dilation_h, dilation_w, channel_per_deformable_group, + batch_size, 2 * kernel_h * kernel_w * deformable_group, deformable_group, height_col, width_col, + grad_offset, grad_mask); + cudaError_t err = cudaGetLastError(); + if (err != cudaSuccess) + { + printf("error in modulated_deformable_col2im_coord_cuda: %s\n", cudaGetErrorString(err)); + } +} \ No newline at end of file diff --git a/centernet/dcnv2Plugin/dcn_v2_im2col_cuda.h b/centernet/dcnv2Plugin/dcn_v2_im2col_cuda.h new file mode 100644 index 0000000..0379b63 --- /dev/null +++ b/centernet/dcnv2Plugin/dcn_v2_im2col_cuda.h @@ -0,0 +1,100 @@ +/*! + ******************* BEGIN Caffe Copyright Notice and Disclaimer **************** + * + * COPYRIGHT + * + * All contributions by the University of California: + * Copyright (c) 2014-2017 The Regents of the University of California (Regents) + * All rights reserved. + * + * All other contributions: + * Copyright (c) 2014-2017, the respective contributors + * All rights reserved. + * + * Caffe uses a shared copyright model: each contributor holds copyright over + * their contributions to Caffe. The project versioning records all such + * contribution and copyright details. If a contributor wants to further mark + * their specific copyright on a particular contribution, they should indicate + * their copyright solely in the commit message of the change when it is + * committed. + * + * LICENSE + * + * Redistribution and use in source and binary forms, with or without + * modification, are permitted provided that the following conditions are met: + * + * 1. Redistributions of source code must retain the above copyright notice, this + * list of conditions and the following disclaimer. + * 2. Redistributions in binary form must reproduce the above copyright notice, + * this list of conditions and the following disclaimer in the documentation + * and/or other materials provided with the distribution. + * + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND + * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED + * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR + * ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES + * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; + * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND + * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT + * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS + * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + * + * CONTRIBUTION AGREEMENT + * + * By contributing to the BVLC/caffe repository through pull-request, comment, + * or otherwise, the contributor releases their content to the + * license and copyright terms herein. + * + ***************** END Caffe Copyright Notice and Disclaimer ******************** + * + * Copyright (c) 2018 Microsoft + * Licensed under The MIT License [see LICENSE for details] + * \file modulated_deformable_im2col.h + * \brief Function definitions of converting an image to + * column matrix based on kernel, padding, dilation, and offset. + * These functions are mainly used in deformable convolution operators. + * \ref: https://arxiv.org/abs/1811.11168 + * \author Yuwen Xiong, Haozhi Qi, Jifeng Dai, Xizhou Zhu, Han Hu + */ + +/***************** Adapted by Charles Shang *********************/ + +#ifndef DCN_V2_IM2COL_CUDA +#define DCN_V2_IM2COL_CUDA + +// #ifdef __cplusplus +// extern "C" +// { +// #endif + + void modulated_deformable_im2col_cuda(cudaStream_t stream, + const float *data_im, const float *data_offset, const float *data_mask, + const int batch_size, const int channels, const int height_im, const int width_im, + const int height_col, const int width_col, const int kernel_h, const int kenerl_w, + const int pad_h, const int pad_w, const int stride_h, const int stride_w, + const int dilation_h, const int dilation_w, + const int deformable_group, float *data_col); + + void modulated_deformable_col2im_cuda(cudaStream_t stream, + const float *data_col, const float *data_offset, const float *data_mask, + const int batch_size, const int channels, const int height_im, const int width_im, + const int height_col, const int width_col, const int kernel_h, const int kenerl_w, + const int pad_h, const int pad_w, const int stride_h, const int stride_w, + const int dilation_h, const int dilation_w, + const int deformable_group, float *grad_im); + + void modulated_deformable_col2im_coord_cuda(cudaStream_t stream, + const float *data_col, const float *data_im, const float *data_offset, const float *data_mask, + const int batch_size, const int channels, const int height_im, const int width_im, + const int height_col, const int width_col, const int kernel_h, const int kenerl_w, + const int pad_h, const int pad_w, const int stride_h, const int stride_w, + const int dilation_h, const int dilation_w, + const int deformable_group, + float *grad_offset, float *grad_mask); + +// #ifdef __cplusplus +// } +// #endif + +#endif \ No newline at end of file diff --git a/centernet/dcnv2Plugin/dcnv2Plugin.cpp b/centernet/dcnv2Plugin/dcnv2Plugin.cpp new file mode 100644 index 0000000..52c1f64 --- /dev/null +++ b/centernet/dcnv2Plugin/dcnv2Plugin.cpp @@ -0,0 +1,402 @@ +#include "dcnv2Plugin.h" +#include + +using namespace nvinfer1; +using nvinfer1::plugin::DeformableConvolutionalLayer; +using nvinfer1::plugin::DCNv2PluginCreator; + +namespace +{ +const char* DCNv2_PLUGIN_VERSION{"1"}; +const char* DCNv2_PLUGIN_NAME{"DCNv2_TRT"}; +} // namespace + +#define CHECK_CUDA(call) \ + do \ + { \ + cudaError_t status = call; \ + if (status != cudaSuccess) \ + { \ + return status; \ + } \ + } while (0) + +PluginFieldCollection DCNv2PluginCreator::mFC{}; +std::vector DCNv2PluginCreator::mPluginAttributes; + +// Parameterized constructor +DeformableConvolutionalLayer::DeformableConvolutionalLayer( + int out_channels, + int kernel, + int deformable_group, + int dilation, + int padding, + int stride, + const Weights* weight, const Weights* bias): + out_channels(out_channels),kernel_size(kernel),deformable_group(deformable_group), + dilation(dilation),padding(padding),stride(stride){ + mWeight = copyToDevice(weight[0].values, weight[0].count); + mBias = copyToDevice(bias[0].values, bias[0].count); +} + +DeformableConvolutionalLayer::DeformableConvolutionalLayer(const void* buffer, size_t length) +{ + const char* d = static_cast(buffer); + const char* a = d; + in_channels = read(d); + height = read(d); + width = read(d); + height_out = read(d); + width_out = read(d); + + out_channels = read(d); + kernel_size = read(d); + deformable_group = read(d); + dilation = read(d); + padding = read(d); + stride = read(d); + + int count = read(d); + mWeight = deserializeToDevice(d, count); + count = read(d); + mBias = deserializeToDevice(d, count); + + ASSERT(d == a + length); +} + +int DeformableConvolutionalLayer::getNbOutputs() const +{ + // Plugin layer has 2 outputs + return 1; +} + +int DeformableConvolutionalLayer::initialize() +{ + size_t oneSize = height_out * width_out * sizeof(float); + std::vector one_((int)oneSize, 1.0f); + CHECK_CUDA(cudaMalloc((void**)&mOne, oneSize)); + CHECK_CUDA(cudaMalloc((void**)&mColumn, in_channels * kernel_size * kernel_size * oneSize)); + CHECK_CUDA(cudaMemcpy(mOne, one_.data(), oneSize, cudaMemcpyHostToDevice)); + return STATUS_SUCCESS; +} + +Dims DeformableConvolutionalLayer::getOutputDimensions(int index, const Dims* inputs, int nbInputs) +{ + ASSERT(index == 0); + ASSERT(nbInputs == 3); + + in_channels = inputs[0].d[0]; + height = inputs[0].d[1]; + width = inputs[0].d[2]; + height_out = (inputs[0].d[1] + 2 * padding - (dilation * (kernel_size - 1) + 1)) / stride + 1; + width_out = (inputs[0].d[2] + 2 * padding - (dilation * (kernel_size - 1) + 1)) / stride + 1; + + return Dims3(out_channels, height_out, width_out); +} + +size_t DeformableConvolutionalLayer::getWorkspaceSize(int maxBatchSize) const +{ + return 0; +} + +int DeformableConvolutionalLayer::enqueue(int batchSize, const void* const* inputs, void** outputs, void* workspace, cudaStream_t stream) +{ + const float* input = static_cast(inputs[0]); + const float* offset = static_cast(inputs[1]); + const float* offset_mask = static_cast(inputs[2]); + const float* mask = offset_mask + deformable_group * 2 * kernel_size * kernel_size * height * width; + float * output = static_cast(outputs[0]); + + float alpha{1}, beta{0}; + + // Do Bias first: + // M,N,K are dims of matrix A and B + // (see http://docs.nvidia.com/cuda/cublas/#cublas-lt-t-gt-gemm) + // (N x 1) (1 x M) + int m_ = out_channels; + int n_ = height_out * width_out; + int k_ = 1; + cublasSgemm(mCublas, CUBLAS_OP_T, CUBLAS_OP_N, n_, m_, k_, &alpha, + mOne, k_, + static_cast(mBias.values), k_, &beta, + output, n_); + + modulated_deformable_im2col_cuda(stream, input, offset, mask, + 1, in_channels, height, width, + height_out, width_out, kernel_size, kernel_size, + padding, padding, stride, stride, dilation, dilation, + deformable_group, mColumn); + + //(k * m) x (m * n) + // Y = WC + int m = out_channels; + int n = height_out * width_out; + int k = in_channels * kernel_size * kernel_size; + cublasSgemm(mCublas, CUBLAS_OP_N, CUBLAS_OP_N, n, m, k, &alpha, + mColumn, n, + static_cast(mWeight.values), k, &alpha, + output, n); + + return 0; +} + +size_t DeformableConvolutionalLayer::getSerializationSize() const +{ + return sizeof(int) * 13 + (mWeight.count + mBias.count) * sizeof(float); +} + +void DeformableConvolutionalLayer::serialize(void* buffer) const +{ + char *d = reinterpret_cast(buffer), *a = d; + write(d, in_channels); + write(d, height); + write(d, width); + write(d, height_out); + write(d, width_out); + + write(d, out_channels); + write(d, kernel_size); + write(d, deformable_group); + write(d, dilation); + write(d, padding); + write(d, stride); + + write(d, (int) mWeight.count); + serializeFromDevice(d, mWeight); + write(d, (int) mBias.count); + serializeFromDevice(d, mBias); + + ASSERT(d == a + getSerializationSize()); +} + +bool DeformableConvolutionalLayer::supportsFormat(DataType type, PluginFormat format) const +{ + return (type == DataType::kFLOAT && format == PluginFormat::kNCHW); +} + +Weights DeformableConvolutionalLayer::copyToDevice(const void* hostData, size_t count) +{ + void* deviceData; + CUASSERT(cudaMalloc(&deviceData, count * sizeof(float))); + CUASSERT(cudaMemcpy(deviceData, hostData, count * sizeof(float), cudaMemcpyHostToDevice)); + return Weights{DataType::kFLOAT, deviceData, int64_t(count)}; +} + +void DeformableConvolutionalLayer::serializeFromDevice(char*& hostBuffer, Weights deviceWeights) const +{ + CUASSERT(cudaMemcpy(hostBuffer, deviceWeights.values, deviceWeights.count * sizeof(float), cudaMemcpyDeviceToHost)); + hostBuffer += deviceWeights.count * sizeof(float); +} + +Weights DeformableConvolutionalLayer::deserializeToDevice(const char*& hostBuffer, size_t count) +{ + Weights w = copyToDevice(hostBuffer, count); + hostBuffer += count * sizeof(float); + return w; +} + +const char* DeformableConvolutionalLayer::getPluginType() const +{ + return DCNv2_PLUGIN_NAME; +} + +const char* DeformableConvolutionalLayer::getPluginVersion() const +{ + return DCNv2_PLUGIN_VERSION; +} + +void DeformableConvolutionalLayer::terminate() { + if (mOne) + { + cudaFree(mOne); + mOne = nullptr; + } + if (mColumn) + { + cudaFree(mColumn); + mColumn = nullptr; + } +} + +void DeformableConvolutionalLayer::destroy() +{ + delete this; +} + +IPluginV2Ext* DeformableConvolutionalLayer::clone() const +{ + IPluginV2Ext* plugin = new DeformableConvolutionalLayer(*this); + plugin->setPluginNamespace(mPluginNamespace.c_str()); + return plugin; +} + +// Set plugin namespace +void DeformableConvolutionalLayer::setPluginNamespace(const char* pluginNamespace) +{ + mPluginNamespace = pluginNamespace; +} + +const char* DeformableConvolutionalLayer::getPluginNamespace() const +{ + return mPluginNamespace.c_str(); +} + +// Return the DataType of the plugin output at the requested index. +DataType DeformableConvolutionalLayer::getOutputDataType(int index, const nvinfer1::DataType* inputTypes, int nbInputs) const +{ + // Only DataType::kFLOAT is acceptable by the plugin layer + return DataType::kFLOAT; +} +// Return true if output tensor is broadcast across a batch. +bool DeformableConvolutionalLayer::isOutputBroadcastAcrossBatch(int outputIndex, const bool* inputIsBroadcasted, int nbInputs) const +{ + return false; +} + +// Return true if plugin can use input that is broadcast across batch without replication. +bool DeformableConvolutionalLayer::canBroadcastInputAcrossBatch(int inputIndex) const +{ + return false; +} + +// Configure the layer with input and output data types. +// inutDims: input Dimensions for the plugin layer +// nInputs : Number of inputs to the plugin layer +// outputDims: output Dimensions from the plugin layer +// nOutputs: number of outputs from the plugin layer +// type: DataType configuration for the plugin layer +// format: format NCHW, NHWC etc +// maxbatchSize: maximum batch size for the plugin layer +void DeformableConvolutionalLayer::configurePlugin(const Dims* inputDims, int nbInputs, const Dims* outputDims, int nbOutputs, + const DataType* inputTypes, const DataType* outputTypes, const bool* inputIsBroadcast, + const bool* outputIsBroadcast, PluginFormat floatFormat, int maxBatchSize) +{ + ASSERT(*inputTypes == DataType::kFLOAT && floatFormat == PluginFormat::kNCHW); +} + +// Attach the plugin object to an execution context and grant the plugin the access to some context resource. +void DeformableConvolutionalLayer::attachToContext(cudnnContext* cudnnContext, cublasContext* cublasContext, IGpuAllocator* gpuAllocator) +{ + mCublas = cublasContext; +} + +// Detach the plugin object from its execution context. +void DeformableConvolutionalLayer::detachFromContext() {} + +DCNv2PluginCreator::DCNv2PluginCreator() +{ + mPluginAttributes.emplace_back(PluginField("out_channels", nullptr, PluginFieldType::kINT32, 1)); + mPluginAttributes.emplace_back(PluginField("kernel", nullptr, PluginFieldType::kINT32, 1)); + mPluginAttributes.emplace_back(PluginField("deformable_group", nullptr, PluginFieldType::kINT32, 1)); + mPluginAttributes.emplace_back(PluginField("dilation", nullptr, PluginFieldType::kINT32, 1)); + mPluginAttributes.emplace_back(PluginField("padding", nullptr, PluginFieldType::kINT32, 1)); + mPluginAttributes.emplace_back(PluginField("stride", nullptr, PluginFieldType::kINT32, 1)); + mPluginAttributes.emplace_back(PluginField("weight", nullptr, PluginFieldType::kFLOAT32, 1)); + mPluginAttributes.emplace_back(PluginField("bias", nullptr, PluginFieldType::kFLOAT32, 1)); + + mFC.nbFields = mPluginAttributes.size(); + mFC.fields = mPluginAttributes.data(); +} + +const char* DCNv2PluginCreator::getPluginName() const +{ + return DCNv2_PLUGIN_NAME; +} + +const char* DCNv2PluginCreator::getPluginVersion() const +{ + return DCNv2_PLUGIN_VERSION; +} + +const PluginFieldCollection* DCNv2PluginCreator::getFieldNames() +{ + return &mFC; +} + +IPluginV2Ext* DCNv2PluginCreator::createPlugin(const char* name, const PluginFieldCollection* fc) +{ + std::vector weight; + std::vector bias; + int out_channels, kernel, deformable_group, padding, stride, dilation; + const PluginField* fields = fc->fields; + for (int i = 0; i < fc->nbFields; ++i) + { + const char* attrName = fields[i].name; + if (!strcmp(attrName, "out_channels")) + { + ASSERT(fields[i].type == PluginFieldType::kINT32); + out_channels = *(static_cast(fields[i].data)); + } + else if (!strcmp(attrName, "kernel")) + { + ASSERT(fields[i].type == PluginFieldType::kINT32); + kernel = *(static_cast(fields[i].data)); + } + else if (!strcmp(attrName, "deformable_group")) + { + ASSERT(fields[i].type == PluginFieldType::kINT32); + deformable_group = *(static_cast(fields[i].data)); + } + else if (!strcmp(attrName, "dilation")) + { + ASSERT(fields[i].type == PluginFieldType::kINT32); + dilation = *(static_cast(fields[i].data)); + } + else if (!strcmp(attrName, "stride")) + { + ASSERT(fields[i].type == PluginFieldType::kINT32); + stride = *(static_cast(fields[i].data)); + } + else if (!strcmp(attrName, "padding")) + { + ASSERT(fields[i].type == PluginFieldType::kINT32); + padding = *(static_cast(fields[i].data)); + } + else if (!strcmp(attrName, "weight")) + { + ASSERT(fields[i].type == PluginFieldType::kFLOAT32); + int size = fields[i].length; + weight.reserve(size); + const auto* w = static_cast(fields[i].data); + for (int j = 0; j < size; j++) + { + weight.push_back(*w); + w++; + } + } + else if (!strcmp(attrName, "bias")) + { + ASSERT(fields[i].type == PluginFieldType::kFLOAT32); + int size = fields[i].length; + bias.reserve(size); + const auto* w = static_cast(fields[i].data); + for (int j = 0; j < size; j++) + { + bias.push_back(*w); + w++; + } + } + } + + Weights mWeight{DataType::kFLOAT, weight.data(), (int64_t) weight.size()}; + Weights mBias{DataType::kFLOAT, bias.data(), (int64_t) bias.size()}; + + DeformableConvolutionalLayer* obj = new DeformableConvolutionalLayer(out_channels, + kernel, + deformable_group, + dilation, + padding, + stride, + &mWeight, &mBias); + obj->setPluginNamespace(mNamespace.c_str()); + return obj; +} + +IPluginV2Ext* DCNv2PluginCreator::deserializePlugin(const char* name, const void* serialData, size_t serialLength) +{ + // This object will be deleted when the network is destroyed, which will + // call Normalize::destroy() + DeformableConvolutionalLayer* obj = new DeformableConvolutionalLayer(serialData, serialLength); + obj->setPluginNamespace(mNamespace.c_str()); + return obj; +} \ No newline at end of file diff --git a/centernet/dcnv2Plugin/dcnv2Plugin.h b/centernet/dcnv2Plugin/dcnv2Plugin.h new file mode 100644 index 0000000..28525b1 --- /dev/null +++ b/centernet/dcnv2Plugin/dcnv2Plugin.h @@ -0,0 +1,152 @@ +/* + * Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#ifndef TRT_DCNV2_PLUGIN_H +#define TRT_DCNV2_PLUGIN_H +#include "kernel.h" +#include "plugin.h" +#include "dcn_v2_im2col_cuda.h" + +#include "serialize.hpp" +#include +#include +#include +#include +#include +#include + +using namespace nvinfer1::plugin; +namespace nvinfer1 +{ +namespace plugin +{ + +class DeformableConvolutionalLayer : public IPluginV2Ext +{ +public: + DeformableConvolutionalLayer(int out_channels, + int kernel, + int deformable_group, + int dilation, + int padding, + int stride, + const Weights* weight, const Weights* bias); + + DeformableConvolutionalLayer(const void* buffer, size_t length); + + ~DeformableConvolutionalLayer() override = default; + + int getNbOutputs() const override; + + Dims getOutputDimensions(int index, const Dims* inputs, int nbInputs) override; + + int initialize() override; + + void terminate() override; + + size_t getWorkspaceSize(int maxBatchSize) const override; + + int enqueue( + int batchSize, const void* const* inputs, void** outputs, void* workspace, cudaStream_t stream) override; + + size_t getSerializationSize() const override; + + void serialize(void* buffer) const override; + + bool supportsFormat(DataType type, PluginFormat format) const override; + + const char* getPluginType() const override; + + const char* getPluginVersion() const override; + + void destroy() override; + + IPluginV2Ext* clone() const override; + + void setPluginNamespace(const char* pluginNamespace) override; + + const char* getPluginNamespace() const override; + + DataType getOutputDataType(int index, const nvinfer1::DataType* inputTypes, int nbInputs) const override; + + bool isOutputBroadcastAcrossBatch(int outputIndex, const bool* inputIsBroadcasted, int nbInputs) const override; + + bool canBroadcastInputAcrossBatch(int inputIndex) const override; + + void attachToContext( + cudnnContext* cudnnContext, cublasContext* cublasContext, IGpuAllocator* gpuAllocator) override; + + void configurePlugin(const Dims* inputDims, int nbInputs, const Dims* outputDims, int nbOutputs, + const DataType* inputTypes, const DataType* outputTypes, const bool* inputIsBroadcast, + const bool* outputIsBroadcast, PluginFormat floatFormat, int maxBatchSize) override; + + void detachFromContext() override; + +private: + Weights copyToDevice(const void* hostData, size_t count); + void serializeFromDevice(char*& hostBuffer, Weights deviceWeights) const; + Weights deserializeToDevice(const char*& hostBuffer, size_t count); + + std::string mPluginNamespace; + + int in_channels{}; + int height_out{}; + int width_out{}; + int height{}; + int width{}; + + int out_channels{}; + int kernel_size{}; + int deformable_group{}; + int dilation{}; + int padding{}; + int stride{}; + + Weights mWeight{}; + Weights mBias{}; + + float* mOne; + float* mColumn; + + cublasHandle_t mCublas; +}; + +class DCNv2PluginCreator : public BaseCreator +{ +public: + DCNv2PluginCreator(); + + ~DCNv2PluginCreator() override = default; + + const char* getPluginName() const override; + + const char* getPluginVersion() const override; + + const PluginFieldCollection* getFieldNames() override; + + IPluginV2Ext* createPlugin(const char* name, const PluginFieldCollection* fc) override; + + IPluginV2Ext* deserializePlugin(const char* name, const void* serialData, size_t serialLength) override; + +private: + static PluginFieldCollection mFC; + + // Parameters for DeformableConvolutionalLayer + static std::vector mPluginAttributes; +}; +} // namespace plugin +} // namespace nvinfer1 + +#endif // TRT_DCNv2_PLUGIN_H diff --git a/centernet/sample/common.py b/centernet/sample/common.py new file mode 100644 index 0000000..0512736 --- /dev/null +++ b/centernet/sample/common.py @@ -0,0 +1,238 @@ +# +# Copyright 1993-2020 NVIDIA Corporation. All rights reserved. +# +# NOTICE TO LICENSEE: +# +# This source code and/or documentation ("Licensed Deliverables") are +# subject to NVIDIA intellectual property rights under U.S. and +# international Copyright laws. +# +# These Licensed Deliverables contained herein is PROPRIETARY and +# CONFIDENTIAL to NVIDIA and is being provided under the terms and +# conditions of a form of NVIDIA software license agreement by and +# between NVIDIA and Licensee ("License Agreement") or electronically +# accepted by Licensee. Notwithstanding any terms or conditions to +# the contrary in the License Agreement, reproduction or disclosure +# of the Licensed Deliverables to any third party without the express +# written consent of NVIDIA is prohibited. +# +# NOTWITHSTANDING ANY TERMS OR CONDITIONS TO THE CONTRARY IN THE +# LICENSE AGREEMENT, NVIDIA MAKES NO REPRESENTATION ABOUT THE +# SUITABILITY OF THESE LICENSED DELIVERABLES FOR ANY PURPOSE. IT IS +# PROVIDED "AS IS" WITHOUT EXPRESS OR IMPLIED WARRANTY OF ANY KIND. +# NVIDIA DISCLAIMS ALL WARRANTIES WITH REGARD TO THESE LICENSED +# DELIVERABLES, INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY, +# NONINFRINGEMENT, AND FITNESS FOR A PARTICULAR PURPOSE. +# NOTWITHSTANDING ANY TERMS OR CONDITIONS TO THE CONTRARY IN THE +# LICENSE AGREEMENT, IN NO EVENT SHALL NVIDIA BE LIABLE FOR ANY +# SPECIAL, INDIRECT, INCIDENTAL, OR CONSEQUENTIAL DAMAGES, OR ANY +# DAMAGES WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, +# WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS +# ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR PERFORMANCE +# OF THESE LICENSED DELIVERABLES. +# +# U.S. Government End Users. These Licensed Deliverables are a +# "commercial item" as that term is defined at 48 C.F.R. 2.101 (OCT +# 1995), consisting of "commercial computer software" and "commercial +# computer software documentation" as such terms are used in 48 +# C.F.R. 12.212 (SEPT 1995) and is provided to the U.S. Government +# only as a commercial end item. Consistent with 48 C.F.R.12.212 and +# 48 C.F.R. 227.7202-1 through 227.7202-4 (JUNE 1995), all +# U.S. Government End Users acquire the Licensed Deliverables with +# only those rights set forth herein. +# +# Any use of the Licensed Deliverables in individual and commercial +# software must include, in the user documentation and internal +# comments to the code, the above Disclaimer and U.S. Government End +# Users Notice. +# + +from itertools import chain +import argparse +import os + +import pycuda.driver as cuda +import pycuda.autoinit +import numpy as np + +import tensorrt as trt + +try: + # Sometimes python2 does not understand FileNotFoundError + FileNotFoundError +except NameError: + FileNotFoundError = IOError + +EXPLICIT_BATCH = 1 << (int)(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH) + +def GiB(val): + return val * 1 << 30 + + +def add_help(description): + parser = argparse.ArgumentParser(description=description, formatter_class=argparse.ArgumentDefaultsHelpFormatter) + args, _ = parser.parse_known_args() + + +def find_sample_data(description="Runs a TensorRT Python sample", subfolder="", find_files=[], err_msg=""): + ''' + Parses sample arguments. + + Args: + description (str): Description of the sample. + subfolder (str): The subfolder containing data relevant to this sample + find_files (str): A list of filenames to find. Each filename will be replaced with an absolute path. + + Returns: + str: Path of data directory. + ''' + + # Standard command-line arguments for all samples. + kDEFAULT_DATA_ROOT = os.path.join(os.sep, "usr", "src", "tensorrt", "data") + parser = argparse.ArgumentParser(description=description, formatter_class=argparse.ArgumentDefaultsHelpFormatter) + parser.add_argument("-d", "--datadir", help="Location of the TensorRT sample data directory, and any additional data directories.", action="append", default=[kDEFAULT_DATA_ROOT]) + args, _ = parser.parse_known_args() + + def get_data_path(data_dir): + # If the subfolder exists, append it to the path, otherwise use the provided path as-is. + data_path = os.path.join(data_dir, subfolder) + if not os.path.exists(data_path): + if data_dir != kDEFAULT_DATA_ROOT: + print("WARNING: " + data_path + " does not exist. Trying " + data_dir + " instead.") + data_path = data_dir + # Make sure data directory exists. + if not (os.path.exists(data_path)) and data_dir != kDEFAULT_DATA_ROOT: + print("WARNING: {:} does not exist. Please provide the correct data path with the -d option.".format(data_path)) + return data_path + + data_paths = [get_data_path(data_dir) for data_dir in args.datadir] + return data_paths, locate_files(data_paths, find_files, err_msg) + +def locate_files(data_paths, filenames, err_msg=""): + """ + Locates the specified files in the specified data directories. + If a file exists in multiple data directories, the first directory is used. + + Args: + data_paths (List[str]): The data directories. + filename (List[str]): The names of the files to find. + + Returns: + List[str]: The absolute paths of the files. + + Raises: + FileNotFoundError if a file could not be located. + """ + found_files = [None] * len(filenames) + for data_path in data_paths: + # Find all requested files. + for index, (found, filename) in enumerate(zip(found_files, filenames)): + if not found: + file_path = os.path.abspath(os.path.join(data_path, filename)) + if os.path.exists(file_path): + found_files[index] = file_path + + # Check that all files were found + for f, filename in zip(found_files, filenames): + if not f or not os.path.exists(f): + raise FileNotFoundError("Could not find {:}. Searched in data paths: {:}\n{:}".format(filename, data_paths, err_msg)) + return found_files + +# Simple helper data class that's a little nicer to use than a 2-tuple. +class HostDeviceMem(object): + def __init__(self, host_mem, device_mem): + self.host = host_mem + self.device = device_mem + + def __str__(self): + return "Host:\n" + str(self.host) + "\nDevice:\n" + str(self.device) + + def __repr__(self): + return self.__str__() + +# Allocates all buffers required for an engine, i.e. host/device inputs/outputs. +def allocate_buffers(engine): + inputs = [] + outputs = [] + bindings = [] + stream = cuda.Stream() + for binding in engine: + size = trt.volume(engine.get_binding_shape(binding)) * engine.max_batch_size + dtype = trt.nptype(engine.get_binding_dtype(binding)) + # Allocate host and device buffers + host_mem = cuda.pagelocked_empty(size, dtype) + device_mem = cuda.mem_alloc(host_mem.nbytes) + # Append the device buffer to device bindings. + bindings.append(int(device_mem)) + # Append to the appropriate list. + if engine.binding_is_input(binding): + inputs.append(HostDeviceMem(host_mem, device_mem)) + else: + outputs.append(HostDeviceMem(host_mem, device_mem)) + return inputs, outputs, bindings, stream + +# This function is generalized for multiple inputs/outputs. +# inputs and outputs are expected to be lists of HostDeviceMem objects. +def do_inference(context, bindings, inputs, outputs, stream, batch_size=1): + # Transfer input data to the GPU. + [cuda.memcpy_htod_async(inp.device, inp.host, stream) for inp in inputs] + # Run inference. + context.execute_async(batch_size=batch_size, bindings=bindings, stream_handle=stream.handle) + # Transfer predictions back from the GPU. + [cuda.memcpy_dtoh_async(out.host, out.device, stream) for out in outputs] + # Synchronize the stream + stream.synchronize() + # Return only the host outputs. + return [out.host for out in outputs] + +# This function is generalized for multiple inputs/outputs for full dimension networks. +# inputs and outputs are expected to be lists of HostDeviceMem objects. +def do_inference_v2(context, bindings, inputs, outputs, stream): + # Transfer input data to the GPU. + [cuda.memcpy_htod_async(inp.device, inp.host, stream) for inp in inputs] + # Run inference. + context.execute_async_v2(bindings=bindings, stream_handle=stream.handle) + # Transfer predictions back from the GPU. + [cuda.memcpy_dtoh_async(out.host, out.device, stream) for out in outputs] + # Synchronize the stream + stream.synchronize() + # Return only the host outputs. + return [out.host for out in outputs] + + +# `retry_call` and `retry` are used to wrap the function we want to try multiple times +def retry_call(func, args=[], kwargs={}, n_retries=3): + """Wrap a function to retry it several times. + + Args: + func: function to call + args (List): args parsed to func + kwargs (Dict): kwargs parsed to func + n_retries (int): maximum times of tries + """ + for i_try in range(n_retries): + try: + func(*args, **kwargs) + break + except: + if i_try == n_retries - 1: + raise + print("retry...") + +# Usage: @retry(n_retries) +def retry(n_retries=3): + """Wrap a function to retry it several times. Decorator version of `retry_call`. + + Args: + n_retries (int): maximum times of tries + + Usage: + @retry(n_retries) + def func(...): + pass + """ + def wrapper(func): + def _wrapper(*args, **kwargs): + retry_call(func, args, kwargs, n_retries) + return _wrapper + return wrapper diff --git a/centernet/sample/test.py b/centernet/sample/test.py new file mode 100644 index 0000000..6e44526 --- /dev/null +++ b/centernet/sample/test.py @@ -0,0 +1,134 @@ +import cv2 as cv +import numpy as np + +import tensorrt as trt +import common + +import torch +import time +from sys import argv + +# You can set the logger severity higher to suppress messages (or lower to display more messages). +TRT_LOGGER = trt.Logger(trt.Logger.WARNING) +trt.init_libnvinfer_plugins(TRT_LOGGER, '') + + +def _gather_feat(feat, ind, mask=None): + dim = feat.size(2) + ind = ind.unsqueeze(2).expand(ind.size(0), ind.size(1), dim) + feat = feat.gather(1, ind) + if mask is not None: + mask = mask.unsqueeze(2).expand_as(feat) + feat = feat[mask] + feat = feat.view(-1, dim) + return feat + + +def _transpose_and_gather_feat(feat, ind): + feat = feat.permute(0, 2, 3, 1).contiguous() + feat = feat.view(feat.size(0), -1, feat.size(3)) + feat = _gather_feat(feat, ind) + return feat + + +def pre_process(image): + long_size = max(image.shape) + img = np.zeros((long_size, long_size, 3)) + img[:image.shape[0], :img.shape[1], :] = image[:] + img = cv.resize(img, (512,512)) + inp_image = ((img / 255. - 0.5) / 0.5).astype(np.float32) + images = inp_image.transpose(2, 0, 1) + return images, long_size/512 + + +def _nms(heat, kernel=3): + pad = (kernel - 1) // 2 + + hmax = torch.nn.functional.max_pool2d( + heat, (kernel, kernel), stride=1, padding=pad) + keep = (hmax == heat).float() + return heat * keep + + +def _topk(scores, K=40): + batch, cat, height, width = scores.size() + + topk_scores, topk_inds = torch.topk(scores.view(batch, cat, -1), K) + + topk_inds = topk_inds % (height * width) + topk_ys = (topk_inds.true_divide(width)).int().float() + topk_xs = (topk_inds % width).int().float() + + topk_score, topk_ind = torch.topk(topk_scores.view(batch, -1), K) + topk_clses = (topk_ind.true_divide(K)).int() + topk_inds = _gather_feat( + topk_inds.view(batch, -1, 1), topk_ind).view(batch, K) + topk_ys = _gather_feat(topk_ys.view(batch, -1, 1), topk_ind).view(batch, K) + topk_xs = _gather_feat(topk_xs.view(batch, -1, 1), topk_ind).view(batch, K) + + return topk_score, topk_inds, topk_clses, topk_ys, topk_xs + + +def ctdet_decode(heat, wh, reg=None, cat_spec_wh=False, K=100): + batch, cat, height, width = heat.size() + + heat = torch.sigmoid(heat) + # perform nms on heatmaps + heat = _nms(heat) + + scores, inds, clses, ys, xs = _topk(heat, K=K) + if reg is not None: + reg = _transpose_and_gather_feat(reg, inds) + reg = reg.view(batch, K, 2) + xs = xs.view(batch, K, 1) + reg[:, :, 0:1] + ys = ys.view(batch, K, 1) + reg[:, :, 1:2] + else: + xs = xs.view(batch, K, 1) + 0.5 + ys = ys.view(batch, K, 1) + 0.5 + wh = _transpose_and_gather_feat(wh, inds) + if cat_spec_wh: + wh = wh.view(batch, K, cat, 2) + clses_ind = clses.view(batch, K, 1, 1).expand(batch, K, 1, 2).long() + wh = wh.gather(2, clses_ind).view(batch, K, 2) + else: + wh = wh.view(batch, K, 2) + clses = clses.view(batch, K, 1).float() + scores = scores.view(batch, K, 1) + bboxes = torch.cat([xs - wh[..., 0:1] / 2, + ys - wh[..., 1:2] / 2, + xs + wh[..., 0:1] / 2, + ys + wh[..., 1:2] / 2], dim=2) + detections = torch.cat([bboxes, scores, clses], dim=2) + return detections + + +if __name__ == '__main__': + try: + engine_path = argv[1] + img_path = argv[2] + except: + print('engine path and image path are needed!') + exit() + with open(engine_path, "rb") as f, trt.Runtime(TRT_LOGGER) as runtime, runtime.deserialize_cuda_engine(f.read()) as engine: + inputs, outputs, bindings, stream = common.allocate_buffers(engine) + with engine.create_execution_context() as context: + img = cv.imread('test.jpg') + dis = img.copy() + img, s = pre_process(img) + # Copy to the pagelocked input buffer + np.copyto(inputs[0].host, img.ravel()) + [hm, wh, reg] = common.do_inference( + context, bindings=bindings, inputs=inputs, outputs=outputs, stream=stream, batch_size=1) + + [dets] = ctdet_decode(torch.from_numpy(hm.reshape(1, 80, 128, 128)), torch.from_numpy( + wh.reshape(1, 2, 128, 128)), torch.from_numpy(reg.reshape(1, 2, 128, 128))) + + for i in dets: + if i[-2] > 0.5: + i[:4] *= 4*s + cv.rectangle(dis, (int(i[0]), int( + i[1])), (int(i[2]), int(i[3])), 255, 1) + cv.putText(dis, '%d' % + int(i[-1]), (int(i[0]), int(i[1])), 1, 1, 255) + + cv.imwrite('trt_out.jpg', dis)