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 <chandler@invix.com>
This commit is contained in:
SsisyphusTao
GitHub
parent
1457e6fcb7
commit
d9e2cbc294
29
centernet/README.md
Normal file
29
centernet/README.md
Normal 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}
|
||||
```
|
||||
|
||||
|
||||
## TODO
|
||||
|
||||
Integrate the post process with trt engine to make it more easier to use.
|
||||
333
centernet/centernet.py
Normal file
333
centernet/centernet.py
Normal 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
21
centernet/dcnv2Plugin/CMakeLists.txt
Normal 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
399
centernet/dcnv2Plugin/dcn_v2_im2col_cuda.cu
Normal 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
100
centernet/dcnv2Plugin/dcn_v2_im2col_cuda.h
Normal 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
402
centernet/dcnv2Plugin/dcnv2Plugin.cpp
Normal 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
152
centernet/dcnv2Plugin/dcnv2Plugin.h
Normal 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
238
centernet/sample/common.py
Normal 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
134
centernet/sample/test.py
Normal 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)
|
||||
Loading…
Reference in New Issue
Block a user