support scale_channels with scale_wh=0

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<int> layers;
         std::string activation = "linear";
 

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

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

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

include/tkDNN/pluginsRT/ScaleChannelsRT.h

@@ -0,0 +1,76 @@
+#include<cassert>
+#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<const dnnType*>(inputs[0]);
+		dnnType *srcDataBack = (dnnType*)reinterpret_cast<const dnnType*>(inputs[1]);
+		dnnType *dstData = reinterpret_cast<dnnType*>(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<char*>(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;
+};

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 "<<layerIdx<<" "<<netLayers[layerIdx]->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));
 

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<int>(buf);
+        bdim.h = readBUF<int>(buf);
+        bdim.w = readBUF<int>(buf);
+        bdim.l = 1;
+
+        ScaleChannelsRT *r = new ScaleChannelsRT(bdim,
+                                                 readBUF<int>(buf)); //scale_wh
+        r->c = readBUF<int>(buf);
+        r->h = readBUF<int>(buf);
+        r->w = readBUF<int>(buf);
+        return r;
+    } 
+
     if(name.find("Pooling") == 0) {
         MaxPoolFixedSizeRT *r = new MaxPoolFixedSizeRT( readBUF<int>(buf), //c
                                                         readBUF<int>(buf), //h

src/ScaleChannels.cpp

@@ -0,0 +1,37 @@
+#include <iostream>
+
+#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;
+}
+
+}}

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 <<<blocks, threads, 0, stream>>>(in_w_h_c, size, channel_size, batch_size, scale_wh, scales_c, out);
+}