diff --git a/include/NetworkRT.h b/include/NetworkRT.h index f4db69a..c1e8d20 100644 --- a/include/NetworkRT.h +++ b/include/NetworkRT.h @@ -45,6 +45,7 @@ public: nvinfer1::ILayer* convert_layer(nvinfer1::ITensor *input, Region *l); nvinfer1::ILayer* convert_layer(nvinfer1::ITensor *input, Shortcut *l); nvinfer1::ILayer* convert_layer(nvinfer1::ITensor *input, Yolo *l); + nvinfer1::ILayer* convert_layer(nvinfer1::ITensor *input, Upsample *l); bool serialize(const char *filename); bool deserialize(const char *filename); diff --git a/src/NetworkRT.cpp b/src/NetworkRT.cpp index 25bfe7c..522f05f 100644 --- a/src/NetworkRT.cpp +++ b/src/NetworkRT.cpp @@ -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(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 bcfb4c7..490fb3c 100644 --- a/src/pluginsRT/RegionRT.cpp +++ b/src/pluginsRT/RegionRT.cpp @@ -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 { diff --git a/src/pluginsRT/UpsampleRT.cpp b/src/pluginsRT/UpsampleRT.cpp new file mode 100644 index 0000000..c7e1e25 --- /dev/null +++ b/src/pluginsRT/UpsampleRT.cpp @@ -0,0 +1,65 @@ +#include +#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(inputs[0]); + dnnType *dstData = reinterpret_cast(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(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; +}; diff --git a/src/pluginsRT/YoloRT.cpp b/src/pluginsRT/YoloRT.cpp index dafbf44..ada4f56 100644 --- a/src/pluginsRT/YoloRT.cpp +++ b/src/pluginsRT/YoloRT.cpp @@ -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 { diff --git a/tests/yolo3_berkeley/yolo3_berkeley.cpp b/tests/yolo3_berkeley/yolo3_berkeley.cpp index 38a863d..2ea1d58 100644 --- a/tests/yolo3_berkeley/yolo3_berkeley.cpp +++ b/tests/yolo3_berkeley/yolo3_berkeley.cpp @@ -77,7 +77,7 @@ const char *c102_bin = "../tests/yolo3_berkeley/layers/c102.bin"; const char *c103_bin = "../tests/yolo3_berkeley/layers/c103.bin"; const char *c104_bin = "../tests/yolo3_berkeley/layers/c104.bin"; const char *c105_bin = "../tests/yolo3_berkeley/layers/c105.bin"; -const char *output_bin = "../tests/yolo3_berkeley/debug/layer82_out.bin"; +const char *output_bin = "../tests/yolo3_berkeley/debug/layer93_out.bin"; int main() { @@ -231,19 +231,21 @@ int main() { tk::dnn::Activation a78 (&net, tk::dnn::ACTIVATION_LEAKY); tk::dnn::Conv2d c79 (&net, 512, 1, 1, 1, 1, 0, 0, c79_bin, true); tk::dnn::Activation a79 (&net, tk::dnn::ACTIVATION_LEAKY); +/* tk::dnn::Conv2d c80 (&net,1024, 3, 3, 1, 1, 1, 1, c80_bin, true); tk::dnn::Activation a80 (&net, tk::dnn::ACTIVATION_LEAKY); tk::dnn::Conv2d c81 (&net, 45, 1, 1, 1, 1, 0, 0, c81_bin, false); - tk::dnn::Yolo g82 (&net, 10, 3); -/* + tk::dnn::Yolo y82 (&net, 10, 3); + tk::dnn::Layer *m83_layers[1] = { &a79 }; tk::dnn::Route m83 (&net, m83_layers, 1); +*/ tk::dnn::Conv2d c84 (&net, 256, 1, 1, 1, 1, 0, 0, c84_bin, true); tk::dnn::Activation a84 (&net, tk::dnn::ACTIVATION_LEAKY); tk::dnn::Upsample u85 (&net, 2); - - tk::dnn::Layer *m86_layers[2] = { &u85, &s61 }; - tk::dnn::Route m86 (&net, m86_layers, 2); + +// tk::dnn::Layer *m86_layers[2] = { &u85, &s61 }; +// tk::dnn::Route m86 (&net, m86_layers, 2); tk::dnn::Conv2d c87 (&net, 256, 1, 1, 1, 1, 0, 0, c87_bin, true); tk::dnn::Activation a87 (&net, tk::dnn::ACTIVATION_LEAKY); tk::dnn::Conv2d c88 (&net, 512, 3, 3, 1, 1, 1, 1, c88_bin, true); @@ -254,19 +256,21 @@ int main() { tk::dnn::Activation a90 (&net, tk::dnn::ACTIVATION_LEAKY); tk::dnn::Conv2d c91 (&net, 256, 1, 1, 1, 1, 0, 0, c91_bin, true); tk::dnn::Activation a91 (&net, tk::dnn::ACTIVATION_LEAKY); +/* tk::dnn::Conv2d c92 (&net, 512, 3, 3, 1, 1, 1, 1, c92_bin, true); tk::dnn::Activation a92 (&net, tk::dnn::ACTIVATION_LEAKY); tk::dnn::Conv2d c93 (&net, 45, 1, 1, 1, 1, 0, 0, c93_bin, false); - tk::dnn::Yolo g94 (&net, 10, 3); + tk::dnn::Yolo y94 (&net, 10, 3); tk::dnn::Layer *m95_layers[1] = { &a91 }; tk::dnn::Route m95 (&net, m95_layers, 1); +*/ tk::dnn::Conv2d c96 (&net, 128, 1, 1, 1, 1, 0, 0, c96_bin, true); tk::dnn::Activation a96 (&net, tk::dnn::ACTIVATION_LEAKY); tk::dnn::Upsample u97 (&net, 2); - tk::dnn::Layer *m98_layers[2] = { &u97, &s36 }; - tk::dnn::Route m98 (&net, m98_layers, 2); +// tk::dnn::Layer *m98_layers[2] = { &u97, &s36 }; +// tk::dnn::Route m98 (&net, m98_layers, 2); tk::dnn::Conv2d c99 (&net, 128, 1, 1, 1, 1, 0, 0, c99_bin, true); tk::dnn::Activation a99 (&net, tk::dnn::ACTIVATION_LEAKY); @@ -282,8 +286,13 @@ int main() { tk::dnn::Conv2d c104 (&net, 256, 3, 3, 1, 1, 1, 1, c104_bin, true); tk::dnn::Activation a104 (&net, tk::dnn::ACTIVATION_LEAKY); tk::dnn::Conv2d c105 (&net, 45, 1, 1, 1, 1, 0, 0, c105_bin, false); - tk::dnn::Yolo g106 (&net, 10, 3); -*/ + tk::dnn::Yolo y106 (&net, 10, 3); + + // merge all yolos +// tk::dnn::Layer *m107_layers[2] = { &y82, &y94, &y106 }; +// tk::dnn::Route m107 (&net, m107_layers, 3); + + // Load input dnnType *data; dnnType *input_h; @@ -318,9 +327,9 @@ int main() { printCenteredTitle(" CHECK RESULTS ", '=', 30); dnnType *out, *out_h; int out_dim = net.getOutputDim().tot(); - readBinaryFile(output_bin, out_dim, &out_h, &out); - std::cout<<"CUDNN vs correct"; checkResult(out_dim, out_data, out); - std::cout<<"TRT vs correct"; checkResult(out_dim, out_data2, out); + //readBinaryFile(output_bin, out_dim, &out_h, &out); + //std::cout<<"CUDNN vs correct"; checkResult(out_dim, out_data, out); + //std::cout<<"TRT vs correct"; checkResult(out_dim, out_data2, out); std::cout<<"CUDNN vs TRT "; checkResult(out_dim, out_data, out_data2); return 0; }