use ElementWiseLayer(broadcast) for scale_channels
This commit is contained in:
@@ -31,7 +31,6 @@ using namespace nvinfer1;
|
|||||||
#include "pluginsRT/RegionRT.h"
|
#include "pluginsRT/RegionRT.h"
|
||||||
#include "pluginsRT/RouteRT.h"
|
#include "pluginsRT/RouteRT.h"
|
||||||
#include "pluginsRT/ShortcutRT.h"
|
#include "pluginsRT/ShortcutRT.h"
|
||||||
#include "pluginsRT/ScaleChannelsRT.h"
|
|
||||||
#include "pluginsRT/YoloRT.h"
|
#include "pluginsRT/YoloRT.h"
|
||||||
#include "pluginsRT/UpsampleRT.h"
|
#include "pluginsRT/UpsampleRT.h"
|
||||||
#include "pluginsRT/ResizeLayerRT.h"
|
#include "pluginsRT/ResizeLayerRT.h"
|
||||||
|
|||||||
@@ -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;
|
|
||||||
};
|
|
||||||
+1
-20
@@ -536,11 +536,7 @@ ILayer* NetworkRT::convert_layer(ITensor *input, Shortcut *l) {
|
|||||||
ILayer* NetworkRT::convert_layer(ITensor *input, ScaleChannels *l) {
|
ILayer* NetworkRT::convert_layer(ITensor *input, ScaleChannels *l) {
|
||||||
ITensor *back_tens = tensors[l->backLayer];
|
ITensor *back_tens = tensors[l->backLayer];
|
||||||
|
|
||||||
IPlugin *plugin = new ScaleChannelsRT(l->backLayer->output_dim, l->scale_wh);
|
IElementWiseLayer *lRT = networkRT->addElementWise(*input, *back_tens, ElementWiseOperation::kPROD);
|
||||||
ITensor **inputs = new ITensor*[2];
|
|
||||||
inputs[0] = input;
|
|
||||||
inputs[1] = back_tens;
|
|
||||||
IPluginLayer *lRT = networkRT->addPlugin(inputs, 2, *plugin);
|
|
||||||
checkNULL(lRT);
|
checkNULL(lRT);
|
||||||
return lRT;
|
return lRT;
|
||||||
}
|
}
|
||||||
@@ -717,21 +713,6 @@ IPlugin* PluginFactory::createPlugin(const char* layerName, const void* serialDa
|
|||||||
return r;
|
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) {
|
if(name.find("Pooling") == 0) {
|
||||||
MaxPoolFixedSizeRT *r = new MaxPoolFixedSizeRT( readBUF<int>(buf), //c
|
MaxPoolFixedSizeRT *r = new MaxPoolFixedSizeRT( readBUF<int>(buf), //c
|
||||||
readBUF<int>(buf), //h
|
readBUF<int>(buf), //h
|
||||||
|
|||||||
Reference in New Issue
Block a user