From 93082cae1f26a48406223a964d3f6542af219312 Mon Sep 17 00:00:00 2001
From: thnkinbtfly <tbvj5914@gmail.com>
Date: Mon, 14 Dec 2020 11:01:23 +0000
Subject: [PATCH] use ElementWiseLayer(broadcast) for scale_channels

---
 include/tkDNN/NetworkRT.h                 |  1 -
 include/tkDNN/pluginsRT/ScaleChannelsRT.h | 76 -----------------------
 src/NetworkRT.cpp                         | 21 +------
 3 files changed, 1 insertion(+), 97 deletions(-)
 delete mode 100644 include/tkDNN/pluginsRT/ScaleChannelsRT.h

diff --git a/include/tkDNN/NetworkRT.h b/include/tkDNN/NetworkRT.h
index 397e4b7..ecb5305 100644
--- a/include/tkDNN/NetworkRT.h
+++ b/include/tkDNN/NetworkRT.h
@@ -31,7 +31,6 @@ 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"
diff --git a/include/tkDNN/pluginsRT/ScaleChannelsRT.h b/include/tkDNN/pluginsRT/ScaleChannelsRT.h
deleted file mode 100644
index 2458d49..0000000
--- a/include/tkDNN/pluginsRT/ScaleChannelsRT.h
+++ /dev/null
@@ -1,76 +0,0 @@
-#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;
-};
diff --git a/src/NetworkRT.cpp b/src/NetworkRT.cpp
index 9721b82..429bf46 100644
--- a/src/NetworkRT.cpp
+++ b/src/NetworkRT.cpp
@@ -536,11 +536,7 @@ 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);
+    IElementWiseLayer *lRT = networkRT->addElementWise(*input, *back_tens, ElementWiseOperation::kPROD);
     checkNULL(lRT);
     return lRT;
 }
@@ -717,21 +713,6 @@ 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