upsample ok, route have problems

This commit is contained in:
Francesco Gatti
2018-12-22 23:56:01 +01:00
parent 34be4cd00f
commit 3bd725801d
6 changed files with 112 additions and 16 deletions
+20
View File
@@ -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;
}