shortcut rt test

This commit is contained in:
Francesco Gatti
2018-12-21 15:53:39 +01:00
parent 7a51b4382d
commit 64626bf547
4 changed files with 97 additions and 3 deletions
+15
View File
@@ -14,6 +14,7 @@ using namespace nvinfer1;
#include "pluginsRT/ActivationLeakyRT.cpp"
#include "pluginsRT/ReorgRT.cpp"
#include "pluginsRT/RegionRT.cpp"
#include "pluginsRT/ShortcutRT.cpp"
#include "pluginsRT/Int8Calibrator.cpp"
// Logger for info/warning/errors
@@ -167,6 +168,8 @@ ILayer* NetworkRT::convert_layer(ITensor *input, Layer *l) {
return convert_layer(input, (Reorg*) l);
if(type == LAYER_REGION)
return convert_layer(input, (Region*) l);
if(type == LAYER_SHORTCUT)
return convert_layer(input, (Shortcut*) l);
FatalError("Layer not implemented in tensorRT");
return NULL;
@@ -321,6 +324,18 @@ ILayer* NetworkRT::convert_layer(ITensor *input, Region *l) {
return lRT;
}
ILayer* NetworkRT::convert_layer(ITensor *input, Shortcut *l) {
//std::cout<<"convert Shortcut\n";
//std::cout<<"New plugin Shortcut\n";
ITensor *tens = tensors[l->backLayer];
IPlugin *plugin = new ShortcutRT(tens);
IPluginLayer *lRT = networkRT->addPlugin(&input, 1, *plugin);
checkNULL(lRT);
return lRT;
}
bool NetworkRT::serialize(const char *filename) {
std::ofstream p(filename);
+64
View File
@@ -0,0 +1,64 @@
#include<cassert>
#include "kernels.h"
class ShortcutRT : public IPlugin {
public:
ShortcutRT(ITensor *tens) {
this->back_layer = tens;
}
~ShortcutRT(){
}
int getNbOutputs() const override {
return 1;
}
Dims getOutputDimensions(int index, const Dims* inputs, int nbInputDims) override {
return DimsCHW{inputs[0].d[0], inputs[0].d[1], inputs[0].d[2]};
}
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 *dstData = reinterpret_cast<dnnType*>(outputs[0]);
checkCuda( cudaMemcpyAsync(dstData, srcData, batchSize*c*h*w*sizeof(dnnType), cudaMemcpyDeviceToDevice, stream));
return 0;
}
virtual size_t getSerializationSize() override {
return 3*sizeof(int);
}
virtual void serialize(void* buffer) override {
char *buf = reinterpret_cast<char*>(buffer);
tk::dnn::writeBUF(buf, c);
tk::dnn::writeBUF(buf, h);
tk::dnn::writeBUF(buf, w);
}
int c, h, w;
ITensor *back_layer;
};