From bec6b236bb1575bcd25dee2867252f6a86d880cb Mon Sep 17 00:00:00 2001 From: thnkinbtfly Date: Mon, 14 Dec 2020 06:10:21 +0000 Subject: [PATCH] support scale_channels with scale_wh=0 --- include/tkDNN/DarknetParser.h | 1 + include/tkDNN/Layer.h | 21 +++++++ include/tkDNN/NetworkRT.h | 2 + include/tkDNN/kernels.h | 4 ++ include/tkDNN/pluginsRT/ScaleChannelsRT.h | 76 +++++++++++++++++++++++ src/DarknetParser.cpp | 14 +++++ src/NetworkRT.cpp | 29 +++++++++ src/ScaleChannels.cpp | 37 +++++++++++ src/kernels/scale_channels.cu | 27 ++++++++ 9 files changed, 211 insertions(+) create mode 100644 include/tkDNN/pluginsRT/ScaleChannelsRT.h create mode 100644 src/ScaleChannels.cpp create mode 100644 src/kernels/scale_channels.cu diff --git a/include/tkDNN/DarknetParser.h b/include/tkDNN/DarknetParser.h index 089c4d6..b90e3f4 100644 --- a/include/tkDNN/DarknetParser.h +++ b/include/tkDNN/DarknetParser.h @@ -28,6 +28,7 @@ namespace tk { namespace dnn { int new_coords= 0; float scale_xy = 1; float nms_thresh = 0.45; + int scale_wh_in_scale_channels = 0; std::vector layers; std::string activation = "linear"; diff --git a/include/tkDNN/Layer.h b/include/tkDNN/Layer.h index 9d6499a..f98d36c 100644 --- a/include/tkDNN/Layer.h +++ b/include/tkDNN/Layer.h @@ -28,6 +28,7 @@ enum layerType_t { LAYER_ROUTE, LAYER_REORG, LAYER_SHORTCUT, + LAYER_SCALECHANNELS, LAYER_UPSAMPLE, LAYER_REGION, LAYER_YOLO @@ -78,6 +79,7 @@ public: case LAYER_ROUTE: return "Route"; case LAYER_REORG: return "Reorg"; case LAYER_SHORTCUT: return "Shortcut"; + case LAYER_SCALECHANNELS: return "ScaleChannels"; case LAYER_UPSAMPLE: return "Upsample"; case LAYER_REGION: return "Region"; case LAYER_YOLO: return "Yolo"; @@ -562,6 +564,25 @@ public: Layer *backLayer; }; +/** + ScaleChannels layer + channelwise-multiplication with another layer +*/ +class ScaleChannels : public Layer { + +public: + ScaleChannels(Network *net, Layer *backLayer, int scale_wh); + virtual ~ScaleChannels(); + virtual layerType_t getLayerType() { return LAYER_SCALECHANNELS; }; + + virtual dnnType* infer(dataDim_t &dim, dnnType* srcData); + +public: + Layer *backLayer; + int scale_wh; +}; + + /** Upsample layer Maintains same dimension but change C*H*W distribution diff --git a/include/tkDNN/NetworkRT.h b/include/tkDNN/NetworkRT.h index 6efff9b..397e4b7 100644 --- a/include/tkDNN/NetworkRT.h +++ b/include/tkDNN/NetworkRT.h @@ -31,6 +31,7 @@ using namespace nvinfer1; #include "pluginsRT/RegionRT.h" #include "pluginsRT/RouteRT.h" #include "pluginsRT/ShortcutRT.h" +#include "pluginsRT/ScaleChannelsRT.h" #include "pluginsRT/YoloRT.h" #include "pluginsRT/UpsampleRT.h" #include "pluginsRT/ResizeLayerRT.h" @@ -109,6 +110,7 @@ public: nvinfer1::ILayer* convert_layer(nvinfer1::ITensor *input, Reorg *l); nvinfer1::ILayer* convert_layer(nvinfer1::ITensor *input, Region *l); nvinfer1::ILayer* convert_layer(nvinfer1::ITensor *input, Shortcut *l); + nvinfer1::ILayer* convert_layer(nvinfer1::ITensor *input, ScaleChannels *l); nvinfer1::ILayer* convert_layer(nvinfer1::ITensor *input, Yolo *l); nvinfer1::ILayer* convert_layer(nvinfer1::ITensor *input, Upsample *l); nvinfer1::ILayer* convert_layer(nvinfer1::ITensor *input, DeformConv2d *l); diff --git a/include/tkDNN/kernels.h b/include/tkDNN/kernels.h index ebb0dbb..19f74bb 100644 --- a/include/tkDNN/kernels.h +++ b/include/tkDNN/kernels.h @@ -28,6 +28,10 @@ void shortcutForward(dnnType *srcData, dnnType *dstData, int n1, int c1, int h1, int n2, int c2, int h2, int w2, int s2, cudaStream_t stream = cudaStream_t(0)); +void scaleChannelsForward(dnnType *in_w_h_c, int size, int channel_size, int batch_size, int scale_wh, + dnnType *scales_c, dnnType *out, + cudaStream_t stream = cudaStream_t(0)); + void upsampleForward(dnnType *srcData, dnnType *dstData, int n, int c, int h, int w, int s, int forward, float scale, cudaStream_t stream = cudaStream_t(0)); diff --git a/include/tkDNN/pluginsRT/ScaleChannelsRT.h b/include/tkDNN/pluginsRT/ScaleChannelsRT.h new file mode 100644 index 0000000..2458d49 --- /dev/null +++ b/include/tkDNN/pluginsRT/ScaleChannelsRT.h @@ -0,0 +1,76 @@ +#include +#include "../kernels.h" + +class ScaleChannelsRT : public IPlugin { + +public: + ScaleChannelsRT(tk::dnn::dataDim_t bdim, int scale_wh) { + this->bc = bdim.c; + this->bh = bdim.h; + this->bw = bdim.w; + this->scale_wh = scale_wh; + } + + ~ScaleChannelsRT(){ + + } + + int getNbOutputs() const override { + return 1; + } + + Dims getOutputDimensions(int index, const Dims* inputs, int nbInputDims) override { + return DimsCHW{bc, bh, bw}; + } + + void configure(const Dims* inputDims, int nbInputs, const Dims* outputDims, int nbOutputs, int maxBatchSize) override { + c = inputDims[0].d[0]; + h = inputDims[0].d[1]; + w = inputDims[0].d[2]; + } + + int initialize() override { + + return 0; + } + + virtual void terminate() override { + } + + virtual size_t getWorkspaceSize(int maxBatchSize) const override { + return 0; + } + + virtual int enqueue(int batchSize, const void*const * inputs, void** outputs, void* workspace, cudaStream_t stream) override { + dnnType *srcData = (dnnType*)reinterpret_cast(inputs[0]); + dnnType *srcDataBack = (dnnType*)reinterpret_cast(inputs[1]); + dnnType *dstData = reinterpret_cast(outputs[0]); + + int size = batchSize * bc * bh * bw; + int channel_size = bh * bw; + int batch_size = bc * bh * bw; + scaleChannelsForward(srcDataBack, size, channel_size, batch_size, scale_wh, srcData, dstData, stream); + + return 0; + } + + + virtual size_t getSerializationSize() override { + return 7*sizeof(int); + } + + virtual void serialize(void* buffer) override { + char *buf = reinterpret_cast(buffer); + tk::dnn::writeBUF(buf, bc); + tk::dnn::writeBUF(buf, bh); + tk::dnn::writeBUF(buf, bw); + tk::dnn::writeBUF(buf, scale_wh); + tk::dnn::writeBUF(buf, c); + tk::dnn::writeBUF(buf, h); + tk::dnn::writeBUF(buf, w); + } + + int c, h, w; + int scale_wh; + int bc, bh, bw; +}; diff --git a/src/DarknetParser.cpp b/src/DarknetParser.cpp index a2441fe..3fda8f0 100644 --- a/src/DarknetParser.cpp +++ b/src/DarknetParser.cpp @@ -80,6 +80,8 @@ namespace tk { namespace dnn { fields.groups = std::stoi(value); else if(name.find("group_id") != std::string::npos) fields.group_id = std::stoi(value); + else if(name.find("scale_wh") != std::string::npos) + fields.scale_wh_in_scale_channels = std::stoi(value); else if(name.find("scale_x_y") != std::string::npos) fields.scale_xy = std::stof(value); else if(name.find("beta_nms") != std::string::npos) @@ -150,6 +152,18 @@ namespace tk { namespace dnn { //std::cout<<"shortcut to "<getLayerName()<<"\n"; netLayers.push_back(new tk::dnn::Shortcut(net, netLayers[layerIdx])); + } else if(f.type == "scale_channels") { + if(f.layers.size() != 1) FatalError("no layers to scale_channels\n"); + int layerIdx = f.layers[0]; + if(layerIdx < 0) + layerIdx = netLayers.size() + layerIdx; + if(layerIdx < 0 || layerIdx >= netLayers.size()) FatalError("impossible to scale_channels\n"); + + int scale_wh = f.scale_wh_in_scale_channels; + if(scale_wh != 0) FatalError("Currently only support scale_wh=0 in scale_channels\n") + + netLayers.push_back(new tk::dnn::ScaleChannels(net, netLayers[layerIdx], scale_wh)); + } else if(f.type == "upsample") { netLayers.push_back(new tk::dnn::Upsample(net, f.stride_x)); diff --git a/src/NetworkRT.cpp b/src/NetworkRT.cpp index 83ddf40..9721b82 100644 --- a/src/NetworkRT.cpp +++ b/src/NetworkRT.cpp @@ -242,6 +242,8 @@ ILayer* NetworkRT::convert_layer(ITensor *input, Layer *l) { return convert_layer(input, (Region*) l); if(type == LAYER_SHORTCUT) return convert_layer(input, (Shortcut*) l); + if(type == LAYER_SCALECHANNELS) + return convert_layer(input, (ScaleChannels*) l); if(type == LAYER_YOLO) return convert_layer(input, (Yolo*) l); if(type == LAYER_UPSAMPLE) @@ -531,6 +533,18 @@ ILayer* NetworkRT::convert_layer(ITensor *input, Shortcut *l) { } } +ILayer* NetworkRT::convert_layer(ITensor *input, ScaleChannels *l) { + ITensor *back_tens = tensors[l->backLayer]; + + IPlugin *plugin = new ScaleChannelsRT(l->backLayer->output_dim, l->scale_wh); + ITensor **inputs = new ITensor*[2]; + inputs[0] = input; + inputs[1] = back_tens; + IPluginLayer *lRT = networkRT->addPlugin(inputs, 2, *plugin); + checkNULL(lRT); + return lRT; +} + ILayer* NetworkRT::convert_layer(ITensor *input, Yolo *l) { //std::cout<<"convert Yolo\n"; @@ -703,6 +717,21 @@ IPlugin* PluginFactory::createPlugin(const char* layerName, const void* serialDa return r; } + if(name.find("ScaleChannels") == 0) { //@note Have to be consistent with string in Layer.h + tk::dnn::dataDim_t bdim; + bdim.c = readBUF(buf); + bdim.h = readBUF(buf); + bdim.w = readBUF(buf); + bdim.l = 1; + + ScaleChannelsRT *r = new ScaleChannelsRT(bdim, + readBUF(buf)); //scale_wh + r->c = readBUF(buf); + r->h = readBUF(buf); + r->w = readBUF(buf); + return r; + } + if(name.find("Pooling") == 0) { MaxPoolFixedSizeRT *r = new MaxPoolFixedSizeRT( readBUF(buf), //c readBUF(buf), //h diff --git a/src/ScaleChannels.cpp b/src/ScaleChannels.cpp new file mode 100644 index 0000000..fe45e21 --- /dev/null +++ b/src/ScaleChannels.cpp @@ -0,0 +1,37 @@ +#include + +#include "Layer.h" +#include "kernels.h" + +namespace tk { namespace dnn { +ScaleChannels::ScaleChannels(Network *net, Layer *backLayer, int scale_wh) : Layer(net) { + + this->backLayer = backLayer; + this->scale_wh = scale_wh; + output_dim = backLayer->output_dim; + checkCuda( cudaMalloc(&dstData, output_dim.tot()*sizeof(dnnType)) ); + + if( backLayer->output_dim.c != input_dim.c ) + FatalError("ScaleChannels dim missmatch"); + +} + +ScaleChannels::~ScaleChannels() { + + checkCuda( cudaFree(dstData) ); +} + +dnnType* ScaleChannels::infer(dataDim_t &dim, dnnType* srcData) { + + int size = output_dim.n * output_dim.c * output_dim.h * output_dim.w; + int channel_size = output_dim.h * output_dim.w; + int batch_size = output_dim.c * output_dim.h * output_dim.w; + scaleChannelsForward(this->backLayer->dstData, size, channel_size, batch_size, scale_wh, srcData, dstData); + + //update data dimensions + dim = output_dim; + + return dstData; +} + +}} \ No newline at end of file diff --git a/src/kernels/scale_channels.cu b/src/kernels/scale_channels.cu new file mode 100644 index 0000000..d00172e --- /dev/null +++ b/src/kernels/scale_channels.cu @@ -0,0 +1,27 @@ +#include "kernels.h" +#include "assert.h" + +// https://github.com/AlexeyAB/darknet/blob/master/src/blas_kernels.cu +__global__ void scale_channels_kernel(float *in_w_h_c, int size, int channel_size, int batch_size, int scale_wh, float *scales_c, float *out) +{ + const int index = blockIdx.x*blockDim.x + threadIdx.x; + if (index < size) { + if (scale_wh) { + int osd_index = index % channel_size + (index / batch_size)*channel_size; + + out[index] = in_w_h_c[index] * scales_c[osd_index]; + } + else { + out[index] = in_w_h_c[index] * scales_c[index / channel_size]; + } + } +} + +void scaleChannelsForward(dnnType *in_w_h_c, int size, int channel_size, int batch_size, int scale_wh, + dnnType *scales_c, dnnType *out, cudaStream_t stream) +{ + int blocks = (size+255)/256; + int threads = 256; + + scale_channels_kernel <<>>(in_w_h_c, size, channel_size, batch_size, scale_wh, scales_c, out); +}