diff --git a/include/Layer.h b/include/Layer.h index a9deccc..0156cf4 100644 --- a/include/Layer.h +++ b/include/Layer.h @@ -267,7 +267,6 @@ public: virtual value_type* infer(dataDim_t &dim, value_type* srcData); -protected: int classes, coords, num; float thresh; diff --git a/include/NetworkRT.h b/include/NetworkRT.h index cb0cae9..98bed30 100644 --- a/include/NetworkRT.h +++ b/include/NetworkRT.h @@ -40,6 +40,8 @@ public: 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); + nvinfer1::ITensor* convert_layer(nvinfer1::ITensor *input, Region *l); + }; diff --git a/src/NetworkRT.cpp b/src/NetworkRT.cpp index bd2794b..1ce105b 100644 --- a/src/NetworkRT.cpp +++ b/src/NetworkRT.cpp @@ -7,6 +7,7 @@ using namespace nvinfer1; #include "pluginsRT/ActivationLeakyRT.cpp" #include "pluginsRT/ReorgRT.cpp" +#include "pluginsRT/RegionRT.cpp" // Logger for info/warning/errors class Logger : public ILogger @@ -111,6 +112,8 @@ ITensor* NetworkRT::convert_layer(ITensor *input, Layer *l) { return convert_layer(input, (Route*) l); if(type == LAYER_REORG) return convert_layer(input, (Reorg*) l); + if(type == LAYER_REGION) + return convert_layer(input, (Region*) l); FatalError("Layer not implemented in tensorRT"); return NULL; @@ -237,4 +240,14 @@ ITensor* NetworkRT::convert_layer(ITensor *input, Reorg *l) { return lRT->getOutput(0); } +ITensor* NetworkRT::convert_layer(ITensor *input, Region *l) { + std::cout<<"convert Region\n"; + + std::cout<<"New plugin REGION\n"; + IPlugin *plugin = new RegionRT(l->classes, l->coords, l->num, l->thresh); + IPluginLayer *lRT = networkRT->addPlugin(&input, 1, *plugin); + checkNULL(lRT); + return lRT->getOutput(0); +} + } \ No newline at end of file diff --git a/src/pluginsRT/RegionRT.cpp b/src/pluginsRT/RegionRT.cpp new file mode 100644 index 0000000..fb5a82d --- /dev/null +++ b/src/pluginsRT/RegionRT.cpp @@ -0,0 +1,89 @@ +#include +#include "kernels.h" + +class RegionRT : public IPlugin { + +public: + RegionRT(int classes, int coords, int num, float thresh) { + + this->classes = classes; + this->coords = coords; + this->num = num; + this->thresh = thresh; + } + + ~RegionRT(){ + + } + + int getNbOutputs() const override { + return 1; + } + + Dims getOutputDimensions(int index, const Dims* inputs, int nbInputDims) override { + return inputs[0]; + } + + 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 { + + value_type *srcData = (value_type*)reinterpret_cast(inputs[0]); + value_type *dstData = reinterpret_cast(outputs[0]); + + checkCuda( cudaMemcpy(dstData, srcData, batchSize*c*h*w*sizeof(value_type), cudaMemcpyDeviceToDevice)); + + for (int b = 0; b < batchSize; ++b){ + for(int n = 0; n < num; ++n){ + int index = entry_index(b, n*w*h, 0, batchSize); + activationLOGISTICForward(srcData + index, dstData + index, 2*w*h); + + index = entry_index(b, n*w*h, coords, batchSize); + activationLOGISTICForward(srcData + index, dstData + index, w*h); + } + } + + //softmax start + int index = entry_index(0, 0, coords + 1, batchSize); + softmaxForward( srcData + index, classes, batchSize*num, + (batchSize*c*h*w)/num, + w*h, 1, w*h, 1, dstData + index); + + return 0; + } + + + virtual size_t getSerializationSize() override { + return 0; + } + + virtual void serialize(void* buffer) override { + } + + int c, h, w; + int classes, coords, num; + float thresh; + + int entry_index(int batch, int location, int entry, int batchSize) { + int n = location / (w*h); + int loc = location % (w*h); + return batch*c*h*w*batchSize + n*w*h*(coords+classes+1) + entry*w*h + loc; + } + +}; diff --git a/tests/yolo/yolo.cpp b/tests/yolo/yolo.cpp index 540b910..99ae44b 100644 --- a/tests/yolo/yolo.cpp +++ b/tests/yolo/yolo.cpp @@ -96,7 +96,7 @@ 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; @@ -126,7 +126,12 @@ int main() { } std::cout<<"\n======= CHECK RESULT =======\n"; - std::cout<<"Diffs: "<