diff --git a/include/Layer.h b/include/Layer.h index ab14a43..a9deccc 100644 --- a/include/Layer.h +++ b/include/Layer.h @@ -251,7 +251,6 @@ public: virtual value_type* infer(dataDim_t &dim, value_type* srcData); -protected: int stride; }; diff --git a/include/NetworkRT.h b/include/NetworkRT.h index 520bf65..cb0cae9 100644 --- a/include/NetworkRT.h +++ b/include/NetworkRT.h @@ -38,6 +38,8 @@ public: nvinfer1::ITensor* convert_layer(nvinfer1::ITensor *input, Dense *l); nvinfer1::ITensor* convert_layer(nvinfer1::ITensor *input, Pooling *l); nvinfer1::ITensor* convert_layer(nvinfer1::ITensor *input, Softmax *l); + nvinfer1::ITensor* convert_layer(nvinfer1::ITensor *input, Route *l); + nvinfer1::ITensor* convert_layer(nvinfer1::ITensor *input, Reorg *l); }; diff --git a/src/NetworkRT.cpp b/src/NetworkRT.cpp index 3551cfc..bd2794b 100644 --- a/src/NetworkRT.cpp +++ b/src/NetworkRT.cpp @@ -1,10 +1,12 @@ #include +#include #include "NvInfer.h" #include "NetworkRT.h" using namespace nvinfer1; #include "pluginsRT/ActivationLeakyRT.cpp" +#include "pluginsRT/ReorgRT.cpp" // Logger for info/warning/errors class Logger : public ILogger @@ -17,6 +19,8 @@ class Logger : public ILogger namespace tkDNN { +std::maptensors; + NetworkRT::NetworkRT(Network *net) { builderRT = createInferBuilder(loggerRT); @@ -33,6 +37,7 @@ NetworkRT::NetworkRT(Network *net) { for(int i=0; inum_layers; i++) { Layer *l = net->layers[i]; input = convert_layer(input, l); + tensors[l] = input; } if(input == NULL) FatalError("conversion failed"); @@ -102,6 +107,10 @@ ITensor* NetworkRT::convert_layer(ITensor *input, Layer *l) { return convert_layer(input, (Activation*) l); if(type == LAYER_SOFTMAX) return convert_layer(input, (Softmax*) l); + if(type == LAYER_ROUTE) + return convert_layer(input, (Route*) l); + if(type == LAYER_REORG) + return convert_layer(input, (Reorg*) l); FatalError("Layer not implemented in tensorRT"); return NULL; @@ -206,4 +215,26 @@ ITensor* NetworkRT::convert_layer(ITensor *input, Softmax *l) { return lRT->getOutput(0); } +ITensor* NetworkRT::convert_layer(ITensor *input, Route *l) { + std::cout<<"convert route\n"; + + ITensor *tens[256]; + for(int i=0; ilayers_n; i++) + tens[i] = tensors[l->layers[i]]; + IConcatenationLayer *lRT = networkRT->addConcatenation(tens, l->layers_n); + checkNULL(lRT); + + return lRT->getOutput(0); +} + +ITensor* NetworkRT::convert_layer(ITensor *input, Reorg *l) { + std::cout<<"convert Reorg\n"; + + std::cout<<"New plugin REORG\n"; + IPlugin *plugin = new ReorgRT(l->stride); + IPluginLayer *lRT = networkRT->addPlugin(&input, 1, *plugin); + checkNULL(lRT); + return lRT->getOutput(0); +} + } \ No newline at end of file diff --git a/src/pluginsRT/ReorgRT.cpp b/src/pluginsRT/ReorgRT.cpp new file mode 100644 index 0000000..9d8014d --- /dev/null +++ b/src/pluginsRT/ReorgRT.cpp @@ -0,0 +1,58 @@ +#include +#include "kernels.h" + +class ReorgRT : public IPlugin { + +public: + ReorgRT(int stride) { + this->stride = stride; + } + + ~ReorgRT(){ + + } + + int getNbOutputs() const override { + return 1; + } + + Dims getOutputDimensions(int index, const Dims* inputs, int nbInputDims) override { + return DimsCHW{inputs[0].d[0]*stride*stride, 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 { + + reorgForward((value_type*)reinterpret_cast(inputs[0]), + reinterpret_cast(outputs[0]), + batchSize, c, h, w, stride); + return 0; + } + + + virtual size_t getSerializationSize() override { + return 0; + } + + virtual void serialize(void* buffer) override { + } + + int c, h, w, stride; +}; diff --git a/src/utils.cpp b/src/utils.cpp index 006fbea..01a2f9e 100644 --- a/src/utils.cpp +++ b/src/utils.cpp @@ -70,7 +70,7 @@ int checkResult(int size, value_type *data_d, value_type *correct_d, bool device if(fabs(data_h[i] - correct_h[i]) > 0.0001) { diffs += 1; if(diffs < 10) - printf("%f %f\n", data_h[i], correct_h[i]); + printf("%d: %f %f\n", i, data_h[i], correct_h[i]); } } diff --git a/tests/yolo/yolo.cpp b/tests/yolo/yolo.cpp index 04d15dc..540b910 100644 --- a/tests/yolo/yolo.cpp +++ b/tests/yolo/yolo.cpp @@ -83,7 +83,7 @@ int main() { tkDNN::Activation a23(&net, tkDNN::ACTIVATION_LEAKY); tkDNN::Conv2d c24(&net, 1024, 3, 3, 1, 1, 1, 1, c24_bin, true); tkDNN::Activation a24(&net, tkDNN::ACTIVATION_LEAKY); -/* + tkDNN::Layer *m25_layers[1] = { &a16 }; tkDNN::Route m25(&net, m25_layers, 1); tkDNN::Conv2d c26(&net, 64, 1, 1, 1, 1, 0, 0, c26_bin, true); @@ -96,9 +96,8 @@ int main() { tkDNN::Conv2d c29(&net, 1024, 3, 3, 1, 1, 1, 1, c29_bin, true); tkDNN::Activation a29(&net, tkDNN::ACTIVATION_LEAKY); tkDNN::Conv2d c30(&net, 425, 1, 1, 1, 1, 0, 0, c30_bin, false); +// tkDNN::Region g31(&net, 80, 4, 5, 0.6f); - tkDNN::Region g31(&net, 80, 4, 5, 0.6f); -*/ // Load input value_type *data; value_type *input_h; @@ -109,7 +108,7 @@ int main() { value_type *out_data, *out_data2; tkDNN::dataDim_t dim1 = dim; - std::cout<<"CUDNN inference:\n"; { + std::cout<<"\n==== CUDNN inference =======\n"; { dim1.print(); //print initial dimension TIMER_START out_data = net.infer(dim1, data); @@ -118,7 +117,7 @@ int main() { } tkDNN::dataDim_t dim2 = dim; - std::cout<<"TENSORRT inference:\n"; { + std::cout<<"\n==== TENSORRT inference ====\n"; { dim2.print(); TIMER_START out_data2 = netRT.infer(dim2, data);