diff --git a/include/NetworkRT.h b/include/NetworkRT.h index 1a4039b..f4db69a 100644 --- a/include/NetworkRT.h +++ b/include/NetworkRT.h @@ -44,6 +44,7 @@ public: nvinfer1::ILayer* convert_layer(nvinfer1::ITensor *input, Reorg *l); 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); bool serialize(const char *filename); bool deserialize(const char *filename); diff --git a/src/NetworkRT.cpp b/src/NetworkRT.cpp index e47aa4f..25bfe7c 100644 --- a/src/NetworkRT.cpp +++ b/src/NetworkRT.cpp @@ -15,6 +15,7 @@ using namespace nvinfer1; #include "pluginsRT/ReorgRT.cpp" #include "pluginsRT/RegionRT.cpp" #include "pluginsRT/ShortcutRT.cpp" +#include "pluginsRT/YoloRT.cpp" #include "pluginsRT/Int8Calibrator.cpp" // Logger for info/warning/errors @@ -170,6 +171,8 @@ ILayer* NetworkRT::convert_layer(ITensor *input, Layer *l) { return convert_layer(input, (Region*) l); if(type == LAYER_SHORTCUT) return convert_layer(input, (Shortcut*) l); + if(type == LAYER_YOLO) + return convert_layer(input, (Yolo*) l); FatalError("Layer not implemented in tensorRT"); return NULL; @@ -337,6 +340,16 @@ ILayer* NetworkRT::convert_layer(ITensor *input, Shortcut *l) { return lRT; } +ILayer* NetworkRT::convert_layer(ITensor *input, Yolo *l) { + //std::cout<<"convert Yolo\n"; + + //std::cout<<"New plugin YOLO\n"; + IPlugin *plugin = new YoloRT(l->classes, l->num); + IPluginLayer *lRT = networkRT->addPlugin(&input, 1, *plugin); + checkNULL(lRT); + return lRT; +} + bool NetworkRT::serialize(const char *filename) { @@ -388,6 +401,23 @@ public: return r; } + if(name.find("Shortcut") == 0) { + ShortcutRT *r = new ShortcutRT(); + r->c = readBUF(buf); + r->h = readBUF(buf); + r->w = readBUF(buf); + return r; + } + + if(name.find("Yolo") == 0) { + YoloRT *r = new YoloRT(readBUF(buf), //classes + readBUF(buf)); //num + 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/YoloRT.cpp b/src/pluginsRT/YoloRT.cpp new file mode 100644 index 0000000..dafbf44 --- /dev/null +++ b/src/pluginsRT/YoloRT.cpp @@ -0,0 +1,86 @@ +#include +#include "kernels.h" + +class YoloRT : public IPlugin { + +public: + YoloRT(int classes, int num) { + + this->classes = classes; + this->num = num; + } + + ~YoloRT(){ + + } + + 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 { + + dnnType *srcData = (dnnType*)reinterpret_cast(inputs[0]); + dnnType *dstData = reinterpret_cast(outputs[0]); + + checkCuda( cudaMemcpyAsync(dstData, srcData, batchSize*c*h*w*sizeof(dnnType), cudaMemcpyDeviceToDevice, stream)); + + 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, stream); + + index = entry_index(b, n*w*h, 4, batchSize); + activationLOGISTICForward(srcData + index, dstData + index, (1+classes)*w*h, stream); + } + } + + return 0; + } + + + virtual size_t getSerializationSize() override { + return 5*sizeof(int) + 1*sizeof(float); + } + + virtual void serialize(void* buffer) override { + char *buf = reinterpret_cast(buffer); + tk::dnn::writeBUF(buf, classes); + tk::dnn::writeBUF(buf, num); + tk::dnn::writeBUF(buf, c); + tk::dnn::writeBUF(buf, h); + tk::dnn::writeBUF(buf, w); + } + + int c, h, w; + int classes, num; + + 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*(4+classes+1) + entry*w*h + loc; + } + +}; diff --git a/tests/yolo3_berkeley/yolo3_berkeley.cpp b/tests/yolo3_berkeley/yolo3_berkeley.cpp index 09782ae..38a863d 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/layer11_out.bin"; +const char *output_bin = "../tests/yolo3_berkeley/debug/layer82_out.bin"; int main() { @@ -106,7 +106,7 @@ int main() { tk::dnn::Conv2d c10 (&net, 128, 3, 3, 1, 1, 1, 1, c10_bin, true); tk::dnn::Activation a10 (&net, tk::dnn::ACTIVATION_LEAKY); tk::dnn::Shortcut s11 (&net, &s8); -/* + tk::dnn::Conv2d c12 (&net, 256, 3, 3, 2, 2, 1, 1, c12_bin, true); tk::dnn::Activation a12 (&net, tk::dnn::ACTIVATION_LEAKY); tk::dnn::Conv2d c13 (&net, 128, 1, 1, 1, 1, 0, 0, c13_bin, true); @@ -235,7 +235,7 @@ int main() { 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::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);