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

use ElementWiseLayer(broadcast) for scale_channels

Browse Source
This commit is contained in:
thnkinbtfly
2020-12-14 11:01:23 +00:00
parent f134fb6230
commit 93082cae1f
3 changed files with 1 additions and 97 deletions
Download Patch File Download Diff File Expand all files Collapse all files
include/tkDNN/NetworkRT.h
-1
View File
@@ -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"
include/tkDNN/pluginsRT/ScaleChannelsRT.h
-76
View File
@@ -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;
};
src/NetworkRT.cpp
+1 -20
View File
@@ -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