upsample ok, route have problems
This commit is contained in:
@@ -16,6 +16,7 @@ using namespace nvinfer1;
|
||||
#include "pluginsRT/RegionRT.cpp"
|
||||
#include "pluginsRT/ShortcutRT.cpp"
|
||||
#include "pluginsRT/YoloRT.cpp"
|
||||
#include "pluginsRT/UpsampleRT.cpp"
|
||||
#include "pluginsRT/Int8Calibrator.cpp"
|
||||
|
||||
// Logger for info/warning/errors
|
||||
@@ -173,6 +174,8 @@ ILayer* NetworkRT::convert_layer(ITensor *input, Layer *l) {
|
||||
return convert_layer(input, (Shortcut*) l);
|
||||
if(type == LAYER_YOLO)
|
||||
return convert_layer(input, (Yolo*) l);
|
||||
if(type == LAYER_UPSAMPLE)
|
||||
return convert_layer(input, (Upsample*) l);
|
||||
|
||||
FatalError("Layer not implemented in tensorRT");
|
||||
return NULL;
|
||||
@@ -350,6 +353,15 @@ ILayer* NetworkRT::convert_layer(ITensor *input, Yolo *l) {
|
||||
return lRT;
|
||||
}
|
||||
|
||||
ILayer* NetworkRT::convert_layer(ITensor *input, Upsample *l) {
|
||||
//std::cout<<"convert Upsample\n";
|
||||
|
||||
//std::cout<<"New plugin UPSAMPLE\n";
|
||||
IPlugin *plugin = new UpsampleRT(l->stride);
|
||||
IPluginLayer *lRT = networkRT->addPlugin(&input, 1, *plugin);
|
||||
checkNULL(lRT);
|
||||
return lRT;
|
||||
}
|
||||
|
||||
bool NetworkRT::serialize(const char *filename) {
|
||||
|
||||
@@ -418,6 +430,14 @@ public:
|
||||
return r;
|
||||
}
|
||||
|
||||
if(name.find("Upsample") == 0) {
|
||||
UpsampleRT *r = new UpsampleRT(readBUF<int>(buf)); //stride
|
||||
r->c = readBUF<int>(buf);
|
||||
r->h = readBUF<int>(buf);
|
||||
r->w = readBUF<int>(buf);
|
||||
return r;
|
||||
}
|
||||
|
||||
FatalError("Cant deserialize Plugin");
|
||||
return NULL;
|
||||
}
|
||||
|
||||
@@ -69,7 +69,7 @@ public:
|
||||
|
||||
|
||||
virtual size_t getSerializationSize() override {
|
||||
return 6*sizeof(int) + 1*sizeof(float);
|
||||
return 6*sizeof(int);
|
||||
}
|
||||
|
||||
virtual void serialize(void* buffer) override {
|
||||
|
||||
@@ -0,0 +1,65 @@
|
||||
#include<cassert>
|
||||
#include "kernels.h"
|
||||
|
||||
class UpsampleRT : public IPlugin {
|
||||
|
||||
public:
|
||||
UpsampleRT(int stride) {
|
||||
this->stride = stride;
|
||||
}
|
||||
|
||||
~UpsampleRT(){
|
||||
|
||||
}
|
||||
|
||||
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]*stride, inputs[0].d[2]*stride);
|
||||
}
|
||||
|
||||
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]);
|
||||
|
||||
fill(dstData, batchSize*c*h*w, 0.0);
|
||||
upsampleForward(srcData, dstData, batchSize, c, h, w, stride, 1, 1);
|
||||
return 0;
|
||||
}
|
||||
|
||||
|
||||
virtual size_t getSerializationSize() override {
|
||||
return 4*sizeof(int);
|
||||
}
|
||||
|
||||
virtual void serialize(void* buffer) override {
|
||||
char *buf = reinterpret_cast<char*>(buffer);
|
||||
tk::dnn::writeBUF(buf, stride);
|
||||
tk::dnn::writeBUF(buf, c);
|
||||
tk::dnn::writeBUF(buf, h);
|
||||
tk::dnn::writeBUF(buf, w);
|
||||
}
|
||||
|
||||
int c, h, w, stride;
|
||||
};
|
||||
@@ -57,12 +57,13 @@ public:
|
||||
}
|
||||
}
|
||||
|
||||
std::cout<<"YOLO END\n";
|
||||
return 0;
|
||||
}
|
||||
|
||||
|
||||
virtual size_t getSerializationSize() override {
|
||||
return 5*sizeof(int) + 1*sizeof(float);
|
||||
return 5*sizeof(int);
|
||||
}
|
||||
|
||||
virtual void serialize(void* buffer) override {
|
||||
|
||||
Reference in New Issue
Block a user