mmr/tkDNN
Archived
1
0
Fork 0
Code Issues 46 Pull Requests 12 Releases 4 Activity

Added constant and zero padding with ISliceLayer,need to add them for tensorrt versions less than 8.2 using IPluginV2Ext instead of ISliceLayer since they dont seem to support reflect and zero

Browse Source
This commit is contained in:
perseusdg
2022-01-04 09:03:04 +05:30
parent ba022663f1
commit b75cecb105
8 changed files with 396 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/NetworkRT.cpp
+30
View File
@@ -456,12 +456,42 @@ ILayer* NetworkRT::convert_layer(ITensor *input, Pooling *l) {
}
ILayer* NetworkRT::convert_layer(ITensor *input,Padding *l){
float rt_ver = float(NV_TENSORRT_MAJOR) +
float(NV_TENSORRT_MINOR)/10 +
float(NV_TENSORRT_PATCH)/100;
#if ((NV_TENSORRT_MAJOR == 8 && NV_TENSORRT_MINOR >= 2) || NV_TENSORRT_MAJOR > 8)
auto *lRT = networkRT->addSlice(*input,Dims3{0,0,0},Dims3{l->output_dim.c,l->output_dim.h,l->output_dim.w},Dims3{0,0,0});
if(l->padding_mode == PADDING_MODE_REFLECTION){
lRT->setMode(SliceMode::kREFLECT);
}else if(l->padding_mode == PADDING_MODE_CONSTANT || l->padding_mode == PADDING_MODE_ZERO){
lRT->setMode(SliceMode::kFILL);
lRT->setInput(4, reinterpret_cast<ITensor &>(l->constant));
}
checkNULL(lRT);
return lRT;
#else
//todo add PADDING_MODE_CONSTANT AND PADDING_MODE_ZERO for tensorrt versions < 8.2
if(l->padding_mode == PADDING_MODE_REFLECTION){
auto creator = getPluginRegistry()->getPluginCreator("ReflectionPaddingRT_tkDNN","1");
std::vector<PluginField> mPluginAttributes;
PluginFieldCollection mFC{};
mPluginAttributes.emplace_back(PluginField("padH",&l->paddingH,PluginFieldType::kINT32,1));
mPluginAttributes.emplace_back(PluginField("padW",&l->paddingW,PluginFieldType::kINT32,1));
mPluginAttributes.emplace_back(PluginField("inputH",&l->input_dim.h,PluginFieldType::kINT32,1));
mPluginAttributes.emplace_back(PluginField("inputW",&l->input_dim.w,PluginFieldType::kINT32,1));
mPluginAttributes.emplace_back(PluginField("outputH",&l->output_dim.h,PluginFieldType::kINT32,1));
mPluginAttributes.emplace_back(PluginField("outputW",&l->output_dim.w,PluginFieldType::kINT32,1));
mPluginAttributes.emplace_back(PluginField("n",&l->input_dim.n,PluginFieldType::kINT32,1));
mFC.nbFields = mPluginAttributes.size();
mFC.fields = mPluginAttributes.data();
auto *plugin = creator->createPlugin(l->getLayerName().c_str(),&mFC);
}
auto *lRT = networkRT->addPluginV2(&input, 1, *plugin);
checkNULL(lRT);
return lRT;
#endif
}
ILayer* NetworkRT::convert_layer(ITensor *input, Activation *l) {
src/Padding.cpp
+11 -1
View File
@@ -7,7 +7,7 @@
#include "kernels.h"
namespace tk{ namespace dnn {
Padding::Padding(Network *net, int32_t pad_h, int32_t pad_w, tkdnnPaddingMode_t padding_mode) : Layer(net) {
Padding::Padding(Network *net, int32_t pad_h, int32_t pad_w, tkdnnPaddingMode_t padding_mode,float constant) : Layer(net) {
this->paddingH = pad_h;
this->paddingW = pad_w;
this->padding_mode = padding_mode;
@@ -15,6 +15,11 @@ namespace tk{ namespace dnn {
output_dim.n = input_dim.n;
output_dim.h = input_dim.h + 2 * (this->paddingH);
output_dim.w = input_dim.w + 2 * (this->paddingW);
if(padding_mode == tkdnnPaddingMode_t::PADDING_MODE_CONSTANT){
this->constant = constant;
}else{
this->constant = 0;
}
checkCuda(cudaMalloc(&dstData,output_dim.tot()*sizeof(dnnType)));
}
@@ -28,6 +33,11 @@ namespace tk{ namespace dnn {
reflection_pad2d_out_forward(paddingH, paddingW, srcData, dstData, input_dim.h, input_dim.w, input_dim.c,
input_dim.n);
}
else if(padding_mode == tkdnnPaddingMode_t::PADDING_MODE_CONSTANT){
constant_pad2d_forward(srcData,dstData,input_dim.h,input_dim.w,output_dim.h,output_dim.w,input_dim.c,
input_dim.n,paddingH,paddingW,constant);
}
dim = output_dim;
return dstData;
}
src/kernels/padding.cu
+43 -1
View File
@@ -2,7 +2,9 @@
#include <thrust/pair.h>
#include <stdio.h>
/*
* Reflection padding is from https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/cuda/ReflectionPad.cu
*/
__device__
inline thrust::pair<int32_t,int32_t> get_index_mapping2d(
int32_t input_dim_x,int32_t input_dim_y,int32_t output_dim_x,
@@ -66,3 +68,43 @@ void reflection_pad2d_out_forward(int32_t pad_h,int32_t pad_w,float *srcData,flo
}
}
/*
* constant padding is inspired from https://github.com/apache/incubator-mxnet/blob/master/src/operator/pad.cu
*/
__global__
void constant_pad2d_kernel(dnnType *srcData,dnnType *dstData,const int32_t padT,const int32_t padL,float constant,int32_t n,int32_t c,int32_t i_h,int32_t i_w,int32_t o_h,int32_t o_w){
int outputPointId = threadIdx.x + blockIdx.x * blockDim.x;
if(outputPointId >= o_h*o_w){
return ;
}
int Ny = i_h;
int Nx = i_w;
int plane = blockIdx.y;
int batch = blockIdx.z;
int outputPointX = outputPointId % o_w;
int outputPointY = outputPointId / o_w;
int checkT = max(0, outputPointY - padT + 1);
int checkB = max(0, padT + Ny - outputPointY);
int checkL = max(0, outputPointX - padL + 1);
int checkR = max(0, padL + Nx - outputPointX);
int inputPointX = min(max(outputPointX - padL, 0), Nx - 1);
int inputPointY = min(max(outputPointY - padT, 0), Ny - 1);
int need_pad = !(checkT * checkB * checkL * checkR);
float value_to_copy = srcData[batch*c*i_h*i_w + plane*i_h*i_w + inputPointY*i_w + inputPointX];
dstData[batch*c*o_w*o_h + plane*o_h*o_w + outputPointY*o_w + outputPointX] = value_to_copy * (!need_pad) + need_pad*constant;
}
void constant_pad2d_forward(dnnType *srcData,dnnType *dstData,int32_t input_h,int32_t input_w,int32_t output_h,
int32_t output_w,int32_t c,int32_t n,int32_t padT,int32_t padL,dnnType constant,cudaStream_t cudaStream){
int32_t output_plane_size = output_h*output_w;
dim3 block_size(output_plane_size>256 ?256:output_plane_size);
dim3 grid_size(ceilDiv(output_plane_size,static_cast<int32_t>(256)),c,n);
constant_pad2d_kernel<<<grid_size,block_size,0,cudaStream>>>(srcData,dstData,padT,padL,constant,n,c,input_h,input_w,output_h,output_w);
}
src/pluginsRT/ReflectionPadding.cpp
+202
View File
@@ -0,0 +1,202 @@
#include <tkDNN/pluginsRT/ReflectionPadding.h>
using namespace nvinfer1;
std::vector<PluginField> ReflectionPaddingRTPluginCreator::mPluginAttributes;
PluginFieldCollection ReflectionPaddingRTPluginCreator::mFC{};
static const char* REFLECTIONPADDINGRT_PLUGIN_VERSION{"1"};
static const char* REFLECTIONPADDINGRT_PLUGIN_NAME{"ReflectionPaddingRT_tkDNN"};
ReflectionPaddingRT::ReflectionPaddingRT(int32_t padH, int32_t padW, int32_t input_h, int32_t input_w, int32_t output_h,
int32_t output_w, int32_t c, int32_t n) {
this->padH = padH;
this->padW = padW;
this->input_h = input_h;
this->input_w = input_w;
this->output_h = output_h;
this->output_w = output_w;
this->n = n;
this->c = c;
}
ReflectionPaddingRT::ReflectionPaddingRT(const void *data, size_t length) {
const char* buf = reinterpret_cast<const char*>(data),*bufcheck=buf;
padH = readBUF<int32_t>(buf);
padW = readBUF<int32_t>(buf);
input_h = readBUF<int32_t>(buf);
input_w = readBUF<int32_t>(buf);
output_h = readBUF<int32_t>(buf);
output_w = readBUF<int32_t>(buf);
n = readBUF<int32_t>(buf);
c = readBUF<int32_t>(buf);
assert(buf = bufcheck + length);
}
ReflectionPaddingRT::~ReflectionPaddingRT() {}
int ReflectionPaddingRT::getNbOutputs() const NOEXCEPT {
return 1;
}
Dims ReflectionPaddingRT::getOutputDimensions(int index, const Dims *inputs, int nbInputDims) NOEXCEPT {
return Dims3{c,output_h,output_w};
}
int ReflectionPaddingRT::initialize() NOEXCEPT {
return 0;
}
void ReflectionPaddingRT::terminate() NOEXCEPT {
}
size_t ReflectionPaddingRT::getWorkspaceSize(int maxBatchSize) const NOEXCEPT {
return 0;
}
#if NV_TENSORRT_MAJOR > 7
int ReflectionPaddingRT::enqueue(int batchSize, const void *const *inputs, void *const *outputs, void *workspace,
cudaStream_t stream) NOEXCEPT {
dnnType* srcData = (dnnType*)reinterpret_cast<const dnnType*>(inputs[0]);
dnnType* dstData = reinterpret_cast<dnnType*>(outputs[0]);
reflection_pad2d_out_forward(padH,padW,srcData,dstData,input_h,input_w,c,n,stream);
return 0;
}
#elif NV_TENSORRT_MAJOR <= 7
int32_t ReflectionPaddingRT::enqueue (int32_t batchSize, const void *const *inputs, void **outputs, void *workspace, cudaStream_t stream){
dnnType* srcData = (dnnType*)reinterpret_cast<const dnnType*>(inputs[0]);
dnnType* dstData = reinterpret_cast<dnnType*>(outputs[0]);
reflection_pad2d_out_forward(padH,padW,srcData,dstData,input_h,input_w,c,n,stream);
return 0;
}
#endif
size_t ReflectionPaddingRT::getSerializationSize() const NOEXCEPT {
return 8*sizeof(int32_t);
}
void ReflectionPaddingRT::serialize(void *buffer) const NOEXCEPT {
char *buf = reinterpret_cast<char*>(buffer),*a=buf;
writeBUF(buf,padH);
writeBUF(buf,padW);
writeBUF(buf,input_h);
writeBUF(buf,input_w);
writeBUF(buf,output_h);
writeBUF(buf,output_w);
writeBUF(buf,n);
writeBUF(buf,c);
}
void ReflectionPaddingRT::destroy() NOEXCEPT {
delete this;
}
const char *ReflectionPaddingRT::getPluginType() const NOEXCEPT {
return REFLECTIONPADDINGRT_PLUGIN_NAME;
}
const char *ReflectionPaddingRT::getPluginVersion() const NOEXCEPT {
return REFLECTIONPADDINGRT_PLUGIN_VERSION;
}
const char *ReflectionPaddingRT::getPluginNamespace() const NOEXCEPT {
return mPluginNamespace.c_str();
}
void ReflectionPaddingRT::setPluginNamespace(const char *pluginNamespace) NOEXCEPT {
mPluginNamespace = pluginNamespace;
}
IPluginV2Ext *ReflectionPaddingRT::clone() const NOEXCEPT {
auto *p = new ReflectionPaddingRT(padH,padW,input_h,input_w,output_h,output_w,c,n);
p->setPluginNamespace(mPluginNamespace.c_str());
return p;
}
DataType
ReflectionPaddingRT::getOutputDataType(int index, const nvinfer1::DataType *inputTypes, int nbInputs) const NOEXCEPT {
return DataType::kFLOAT;
}
void ReflectionPaddingRT::attachToContext(cudnnContext *cudnnContext, cublasContext *cublasContext,
IGpuAllocator *gpuAllocator) NOEXCEPT {
}
bool ReflectionPaddingRT::isOutputBroadcastAcrossBatch(int outputIndex, const bool *inputIsBroadcasted,
int nbInputs) const NOEXCEPT {
return false;
}
bool ReflectionPaddingRT::canBroadcastInputAcrossBatch(int inputIndex) const NOEXCEPT {
return false;
}
void
ReflectionPaddingRT::configurePlugin(const Dims *inputDims, int32_t nbInputs, const Dims *outputDims, int32_t nbOutputs,
const DataType *inputTypes, const DataType *outputTypes,
const bool *inputIsBroadcast, const bool *outputIsBroadcast,
PluginFormat floatFormat, int32_t maxBatchSize) NOEXCEPT {
}
void ReflectionPaddingRT::detachFromContext() NOEXCEPT {
}
bool ReflectionPaddingRT::supportsFormat(DataType type, PluginFormat format) const NOEXCEPT {
return (type == DataType::kFLOAT && format == PluginFormat::kLINEAR);
}
ReflectionPaddingRTPluginCreator::ReflectionPaddingRTPluginCreator() {
mPluginAttributes.clear();
mFC.nbFields = mPluginAttributes.size();
mFC.fields = mPluginAttributes.data();
}
void ReflectionPaddingRTPluginCreator::setPluginNamespace(const char *pluginNamespace) NOEXCEPT {
mPluginNamespace = pluginNamespace;
}
const char *ReflectionPaddingRTPluginCreator::getPluginNamespace() const NOEXCEPT {
return mPluginNamespace.c_str();
}
IPluginV2Ext *ReflectionPaddingRTPluginCreator::deserializePlugin(const char *name, const void *serialData,
size_t serialLength) NOEXCEPT {
auto *pluginObj = new ReflectionPaddingRT(serialData,serialLength);
pluginObj->setPluginNamespace(mPluginNamespace.c_str());
return pluginObj;
}
IPluginV2Ext *
ReflectionPaddingRTPluginCreator::createPlugin(const char *name, const PluginFieldCollection *fc) NOEXCEPT {
const PluginField *fields = fc->fields;
int padH = *(static_cast<const int32_t*>(fields[0].data));
int padW = *(static_cast<const int32_t*>(fields[1].data));
int inputH = *(static_cast<const int32_t*>(fields[2].data));
int inputW = *(static_cast<const int32_t*>(fields[3].data));
int outputH = *(static_cast<const int32_t*>(fields[4].data));
int outputW = *(static_cast<const int32_t*>(fields[5].data));
int n = *(static_cast<const int32_t*>(fields[6].data));
int c = *(static_cast<const int32_t*>(fields[7].data));
auto *pluginObj = new ReflectionPaddingRT(padH,padW,inputH,inputW,outputH,outputW,c,n);
pluginObj->setPluginNamespace(mPluginNamespace.c_str());
return pluginObj;
}
const char *ReflectionPaddingRTPluginCreator::getPluginName() const NOEXCEPT {
return REFLECTIONPADDINGRT_PLUGIN_NAME;
}
const char *ReflectionPaddingRTPluginCreator::getPluginVersion() const NOEXCEPT {
return REFLECTIONPADDINGRT_PLUGIN_VERSION;
}
const PluginFieldCollection *ReflectionPaddingRTPluginCreator::getFieldNames() NOEXCEPT {
return &mFC;
}