From 3b2f062dd9e75aee0d963c3335133de137c04c16 Mon Sep 17 00:00:00 2001 From: Francesco Gatti Date: Fri, 11 Aug 2017 17:17:05 +0200 Subject: [PATCH] tensorRT serialization OK --- include/NetworkRT.h | 2 +- src/NetworkRT.cpp | 43 ++++++++++++++++++++++++++++------- src/pluginsRT/RegionRT.cpp | 10 +++++++- src/pluginsRT/ReorgRT.cpp | 7 +++++- tests/mnist/test_mnist.cpp | 2 +- tests/yolo-tiny/yolo-tiny.cpp | 2 +- tests/yolo/yolo.cpp | 2 +- 7 files changed, 54 insertions(+), 14 deletions(-) diff --git a/include/NetworkRT.h b/include/NetworkRT.h index 4b6ff73..0a57e2a 100644 --- a/include/NetworkRT.h +++ b/include/NetworkRT.h @@ -25,7 +25,7 @@ public: dnnType *output; cudaStream_t stream; - NetworkRT(Network *net); + NetworkRT(Network *net, const char *name); virtual ~NetworkRT(); /** diff --git a/src/NetworkRT.cpp b/src/NetworkRT.cpp index 9ea82cb..df050b7 100644 --- a/src/NetworkRT.cpp +++ b/src/NetworkRT.cpp @@ -1,5 +1,7 @@ #include #include +#include + #include "NvInfer.h" #include "NetworkRT.h" @@ -22,7 +24,7 @@ namespace tkDNN { std::maptensors; -NetworkRT::NetworkRT(Network *net) { +NetworkRT::NetworkRT(Network *net, const char *name) { float rt_ver = float(NV_TENSORRT_MAJOR) + float(NV_TENSORRT_MINOR)/10 + @@ -35,7 +37,7 @@ NetworkRT::NetworkRT(Network *net) { //add input layer dataDim_t dim = net->layers[0]->input_dim; - if(!fileExist("net.rt")) { + if(!fileExist(name)) { ITensor *input = networkRT->addInput("data", dtRT, DimsCHW{ dim.c, dim.h, dim.w}); checkNULL(input); @@ -64,9 +66,9 @@ NetworkRT::NetworkRT(Network *net) { engineRT = builderRT->buildCudaEngine(*networkRT); // we don't need the network any more //networkRT->destroy(); - serialize("net.rt"); + serialize(name); } else { - deserialize("net.rt"); + deserialize(name); } std::cout<<"create execution context\n"; @@ -220,11 +222,16 @@ ILayer* NetworkRT::convert_layer(ITensor *input, Activation *l) { IPluginLayer *lRT = networkRT->addPlugin(&input, 1, *plugin); checkNULL(lRT); return lRT; - } - IActivationLayer *lRT = networkRT->addActivation(*input, ActivationType::kRELU); - checkNULL(lRT); - return lRT; + } else if(l->act_mode == CUDNN_ACTIVATION_RELU) { + IActivationLayer *lRT = networkRT->addActivation(*input, ActivationType::kRELU); + checkNULL(lRT); + return lRT; + + } else { + FatalError("this Activation mode is not yet implemented"); + return NULL; + } } ILayer* NetworkRT::convert_layer(ITensor *input, Softmax *l) { @@ -297,7 +304,27 @@ public: ActivationLeakyRT *a = new ActivationLeakyRT(); a->size = readBUF(buf); return a; + } + + if(name.find("Region") == 0) { + RegionRT *r = new RegionRT(readBUF(buf), //classes + readBUF(buf), //coords + readBUF(buf), //num + readBUF(buf)); //thesh + r->c = readBUF(buf); + r->h = readBUF(buf); + r->w = readBUF(buf); + return r; } + + if(name.find("Reorg") == 0) { + ReorgRT *r = new ReorgRT(readBUF(buf)); //stride + r->c = readBUF(buf); + r->h = readBUF(buf); + r->w = readBUF(buf); + return r; + } + FatalError("Cant deserialize Plugin"); return NULL; } diff --git a/src/pluginsRT/RegionRT.cpp b/src/pluginsRT/RegionRT.cpp index 2974d87..880701b 100644 --- a/src/pluginsRT/RegionRT.cpp +++ b/src/pluginsRT/RegionRT.cpp @@ -70,10 +70,18 @@ public: virtual size_t getSerializationSize() override { - return 0; + return 6*sizeof(int) + 1*sizeof(float); } virtual void serialize(void* buffer) override { + char *buf = reinterpret_cast(buffer); + tkDNN::writeBUF(buf, classes); + tkDNN::writeBUF(buf, coords); + tkDNN::writeBUF(buf, num); + tkDNN::writeBUF(buf, thresh); + tkDNN::writeBUF(buf, c); + tkDNN::writeBUF(buf, h); + tkDNN::writeBUF(buf, w); } int c, h, w; diff --git a/src/pluginsRT/ReorgRT.cpp b/src/pluginsRT/ReorgRT.cpp index 698101b..d85b366 100644 --- a/src/pluginsRT/ReorgRT.cpp +++ b/src/pluginsRT/ReorgRT.cpp @@ -48,10 +48,15 @@ public: virtual size_t getSerializationSize() override { - return 0; + return 4*sizeof(int); } virtual void serialize(void* buffer) override { + char *buf = reinterpret_cast(buffer); + tkDNN::writeBUF(buf, stride); + tkDNN::writeBUF(buf, c); + tkDNN::writeBUF(buf, h); + tkDNN::writeBUF(buf, w); } int c, h, w, stride; diff --git a/tests/mnist/test_mnist.cpp b/tests/mnist/test_mnist.cpp index 747b2d7..49c74a4 100644 --- a/tests/mnist/test_mnist.cpp +++ b/tests/mnist/test_mnist.cpp @@ -22,7 +22,7 @@ int main() { tkDNN::Dense l6(&net, 10, d3_bin); tkDNN::Softmax l7(&net); - tkDNN::NetworkRT netRT(&net); + tkDNN::NetworkRT netRT(&net, "mnist.rt"); // Load input dnnType *data; diff --git a/tests/yolo-tiny/yolo-tiny.cpp b/tests/yolo-tiny/yolo-tiny.cpp index 0139f78..954fe5f 100644 --- a/tests/yolo-tiny/yolo-tiny.cpp +++ b/tests/yolo-tiny/yolo-tiny.cpp @@ -60,7 +60,7 @@ int main() { net.print(); //convert network to tensorRT - tkDNN::NetworkRT netRT(&net); + tkDNN::NetworkRT netRT(&net, "yolo-tiny.rt"); dnnType *out_data, *out_data2; // cudnn output, tensorRT output diff --git a/tests/yolo/yolo.cpp b/tests/yolo/yolo.cpp index 8bd6e09..50aec86 100644 --- a/tests/yolo/yolo.cpp +++ b/tests/yolo/yolo.cpp @@ -108,7 +108,7 @@ int main() { net.print(); //convert network to tensorRT - tkDNN::NetworkRT netRT(&net); + tkDNN::NetworkRT netRT(&net, "yolo.rt"); dnnType *out_data, *out_data2; // cudnn output, tensorRT output