haotian/duan8
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

add centernet dla34 ctdet task. (#558)

Browse Source 
* add centernet dla34 ctdet task.

* update readme.

* update readme and fix a bug in sample.

Co-authored-by: chandler <chandler@invix.com>
This commit is contained in:
SsisyphusTao 2021-05-23 16:29:14 +08:00 committed by GitHub
parent 1457e6fcb7
commit d9e2cbc294
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
9 changed files with 1808 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

29
centernet/README.md Normal file
View File

@ -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.

333
centernet/centernet.py Normal file
View File

@ -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)

21
centernet/dcnv2Plugin/CMakeLists.txt Normal file
View File

@ -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)

399
centernet/dcnv2Plugin/dcn_v2_im2col_cuda.cu Normal file
View File

@ -0,0 +1,399 @@
#include "dcn_v2_im2col_cuda.h"
#include <cstdio>
#include <algorithm>
#include <cstring>
#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<float>(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
<<<GET_BLOCKS(num_kernels), CUDA_NUM_THREADS,
0, stream>>>(
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
<<<GET_BLOCKS(num_kernels), CUDA_NUM_THREADS,
0, stream>>>(
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
<<<GET_BLOCKS(num_kernels), CUDA_NUM_THREADS,
0, stream>>>(
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));
}
}

100
centernet/dcnv2Plugin/dcn_v2_im2col_cuda.h Normal file
View File

@ -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

402
centernet/dcnv2Plugin/dcnv2Plugin.cpp Normal file
View File

@ -0,0 +1,402 @@
#include "dcnv2Plugin.h"
#include <iostream>
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<PluginField> 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<const char*>(buffer);
const char* a = d;
in_channels = read<int>(d);
height = read<int>(d);
width = read<int>(d);
height_out = read<int>(d);
width_out = read<int>(d);
out_channels = read<int>(d);
kernel_size = read<int>(d);
deformable_group = read<int>(d);
dilation = read<int>(d);
padding = read<int>(d);
stride = read<int>(d);
int count = read<int>(d);
mWeight = deserializeToDevice(d, count);
count = read<int>(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<float> 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<const float *>(inputs[0]);
const float* offset = static_cast<const float *>(inputs[1]);
const float* offset_mask = static_cast<const float *>(inputs[2]);
const float* mask = offset_mask + deformable_group * 2 * kernel_size * kernel_size * height * width;
float * output = static_cast<float *>(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<const float *>(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<const float *>(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<char*>(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<float> weight;
std::vector<float> 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<const int*>(fields[i].data));
}
else if (!strcmp(attrName, "kernel"))
{
ASSERT(fields[i].type == PluginFieldType::kINT32);
kernel = *(static_cast<const int*>(fields[i].data));
}
else if (!strcmp(attrName, "deformable_group"))
{
ASSERT(fields[i].type == PluginFieldType::kINT32);
deformable_group = *(static_cast<const int*>(fields[i].data));
}
else if (!strcmp(attrName, "dilation"))
{
ASSERT(fields[i].type == PluginFieldType::kINT32);
dilation = *(static_cast<const int*>(fields[i].data));
}
else if (!strcmp(attrName, "stride"))
{
ASSERT(fields[i].type == PluginFieldType::kINT32);
stride = *(static_cast<const int*>(fields[i].data));
}
else if (!strcmp(attrName, "padding"))
{
ASSERT(fields[i].type == PluginFieldType::kINT32);
padding = *(static_cast<const int*>(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<const float*>(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<const float*>(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;
}

152
centernet/dcnv2Plugin/dcnv2Plugin.h Normal file
View File

@ -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 <cudnn.h>
#include <vector>
#include <cublas_v2.h>
#include <cuda.h>
#include <string>
#include <vector>
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<PluginField> mPluginAttributes;
};
} // namespace plugin
} // namespace nvinfer1
#endif // TRT_DCNv2_PLUGIN_H

238
centernet/sample/common.py Normal file
View File

@ -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

134
centernet/sample/test.py Normal file
View File

@ -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)