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;
+};