diff --git a/demo/demo/demo.cpp b/demo/demo/demo.cpp index a2ea5b1..fbdb2ff 100644 --- a/demo/demo/demo.cpp +++ b/demo/demo/demo.cpp @@ -24,7 +24,7 @@ int main(int argc, char *argv[]) { signal(SIGINT, sig_handler); tk::dnn::Yolo3Detection yolo; - yolo.init("yolo3_berkeley"); + yolo.init("yolo3_berkeley.rt"); gRun = true; diff --git a/include/Layer.h b/include/Layer.h index 860d73a..c7871b5 100644 --- a/include/Layer.h +++ b/include/Layer.h @@ -343,7 +343,7 @@ public: float x, y, w, h; }; - typedef struct detection{ + struct detection{ Yolo::box bbox; int classes; float *prob; @@ -361,7 +361,7 @@ public: dnnType *bias_h, *bias_d; //anchors virtual dnnType* infer(dataDim_t &dim, dnnType* srcData); - int computeDetections(Yolo::detection *dets, int &ndets, int w, int h, int netw, int neth, float thresh); + int computeDetections(Yolo::detection *dets, int &ndets, int netw, int neth, float thresh); dnnType *predictions; diff --git a/include/Network.h b/include/Network.h index 958129c..dc2b5b3 100644 --- a/include/Network.h +++ b/include/Network.h @@ -15,7 +15,7 @@ namespace tk { namespace dnn { */ struct dataDim_t { - int n, c, h, w, l; + int n = 0, c = 0, h = 0, w = 0, l = 0; dataDim_t() : n(1), c(1), h(1), w(1), l(1) {}; diff --git a/include/NetworkRT.h b/include/NetworkRT.h index 5a11a32..d30da03 100644 --- a/include/NetworkRT.h +++ b/include/NetworkRT.h @@ -1,6 +1,7 @@ #ifndef NETWORKRT_H #define NETWORKRT_H +#include // memcpy #include "utils.h" #include "Network.h" #include "Layer.h" @@ -8,6 +9,40 @@ namespace tk { namespace dnn { +template void writeBUF(char*& buffer, const T& val) +{ + *reinterpret_cast(buffer) = val; + buffer += sizeof(T); +} + +template T readBUF(const char*& buffer) +{ + T val = *reinterpret_cast(buffer); + buffer += sizeof(T); + return val; +} + +using namespace nvinfer1; +#include "pluginsRT/ActivationLeakyRT.h" +#include "pluginsRT/ReorgRT.h" +#include "pluginsRT/RegionRT.h" +//#include "pluginsRT/RouteRT.h" +#include "pluginsRT/ShortcutRT.h" +#include "pluginsRT/YoloRT.h" +#include "pluginsRT/UpsampleRT.h" +//#include "pluginsRT/Int8Calibrator.h" + +class PluginFactory : IPluginFactory +{ +public: + YoloRT *yolos[16]; + int n_yolos; + + virtual IPlugin* createPlugin(const char* layerName, const void* serialData, size_t serialLength); +}; + + + class NetworkRT { public: @@ -27,6 +62,8 @@ public: dnnType *output; cudaStream_t stream; + PluginFactory *pluginFactory; + NetworkRT(Network *net, const char *name); virtual ~NetworkRT(); @@ -53,19 +90,5 @@ public: bool deserialize(const char *filename); }; - -template void writeBUF(char*& buffer, const T& val) -{ - *reinterpret_cast(buffer) = val; - buffer += sizeof(T); -} - -template T readBUF(const char*& buffer) -{ - T val = *reinterpret_cast(buffer); - buffer += sizeof(T); - return val; -} - }} #endif //NETWORKRT_H diff --git a/include/Yolo3Detection.h b/include/Yolo3Detection.h index 4cac01d..0ff4b1e 100644 --- a/include/Yolo3Detection.h +++ b/include/Yolo3Detection.h @@ -31,8 +31,8 @@ class Yolo3Detection { cv::Mat bgr[3]; public: - int classes = 10; - int num = 3; + int classes = 0; + int num = 0; float thresh = 0.3; cv::Scalar colors[256]; diff --git a/src/pluginsRT/ActivationLeakyRT.cpp b/include/pluginsRT/ActivationLeakyRT.h similarity index 100% rename from src/pluginsRT/ActivationLeakyRT.cpp rename to include/pluginsRT/ActivationLeakyRT.h diff --git a/src/pluginsRT/Int8Calibrator.cpp b/include/pluginsRT/Int8Calibrator.h similarity index 100% rename from src/pluginsRT/Int8Calibrator.cpp rename to include/pluginsRT/Int8Calibrator.h diff --git a/src/pluginsRT/RegionRT.cpp b/include/pluginsRT/RegionRT.h similarity index 100% rename from src/pluginsRT/RegionRT.cpp rename to include/pluginsRT/RegionRT.h diff --git a/src/pluginsRT/ReorgRT.cpp b/include/pluginsRT/ReorgRT.h similarity index 100% rename from src/pluginsRT/ReorgRT.cpp rename to include/pluginsRT/ReorgRT.h diff --git a/src/pluginsRT/RouteRT.cpp b/include/pluginsRT/RouteRT.h similarity index 100% rename from src/pluginsRT/RouteRT.cpp rename to include/pluginsRT/RouteRT.h diff --git a/src/pluginsRT/ShortcutRT.cpp b/include/pluginsRT/ShortcutRT.h similarity index 100% rename from src/pluginsRT/ShortcutRT.cpp rename to include/pluginsRT/ShortcutRT.h diff --git a/src/pluginsRT/UpsampleRT.cpp b/include/pluginsRT/UpsampleRT.h similarity index 100% rename from src/pluginsRT/UpsampleRT.cpp rename to include/pluginsRT/UpsampleRT.h diff --git a/src/pluginsRT/YoloRT.cpp b/include/pluginsRT/YoloRT.h similarity index 100% rename from src/pluginsRT/YoloRT.cpp rename to include/pluginsRT/YoloRT.h diff --git a/src/NetworkRT.cpp b/src/NetworkRT.cpp index a9e4bc8..9b78e99 100644 --- a/src/NetworkRT.cpp +++ b/src/NetworkRT.cpp @@ -11,14 +11,6 @@ #include "NetworkRT.h" using namespace nvinfer1; -#include "pluginsRT/ActivationLeakyRT.cpp" -#include "pluginsRT/ReorgRT.cpp" -#include "pluginsRT/RegionRT.cpp" -//#include "pluginsRT/RouteRT.cpp" -#include "pluginsRT/ShortcutRT.cpp" -#include "pluginsRT/YoloRT.cpp" -#include "pluginsRT/UpsampleRT.cpp" -#include "pluginsRT/Int8Calibrator.cpp" // Logger for info/warning/errors class Logger : public ILogger { @@ -54,14 +46,7 @@ NetworkRT::NetworkRT(Network *net, const char *name) { builderRT->setMaxBatchSize(1); builderRT->setMaxWorkspaceSize(1 << 30); - //change datatype based on system specs - if(builderRT->platformHasFastInt8()) { - BatchStream bstream({32,dim.c, dim.h, dim.w}, 32, 1); - Int8EntropyCalibrator calib(bstream, 0, false); - builderRT->setInt8Mode(true); - builderRT->setInt8Calibrator(&calib); - - } else if(net->fp16 && builderRT->platformHasFastFp16()) { + if(net->fp16 && builderRT->platformHasFastFp16()) { dtRT = DataType::kHALF; builderRT->setHalf2Mode(true); } @@ -393,87 +378,6 @@ bool NetworkRT::serialize(const char *filename) { return true; } -class PluginFactory : IPluginFactory -{ -public: - - virtual IPlugin* createPlugin(const char* layerName, const void* serialData, size_t serialLength) { - const char * buf = reinterpret_cast(serialData); - - std::string name(layerName); - - if(name.find("Activation") == 0) { - 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 - - 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; - } - - 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); - for(int i=0; inum; i++) - r->mask[i] = readBUF(buf); - for(int i=0; i<3*2*r->num; i++) - r->bias[i] = readBUF(buf); - - std::cout<<"YOLO: "<c<<" "<h<<" "<w<<"\n"; - 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; - } -/* - if(name.find("Route") == 0) { - RouteRT *r = new RouteRT(); - r->in = readBUF(buf); - for(int i=0; ic_in[i] = readBUF(buf); - r->c = readBUF(buf); - r->h = readBUF(buf); - r->w = readBUF(buf); - return r; - } -*/ - FatalError("Cant deserialize Plugin"); - return NULL; - } -}; - bool NetworkRT::deserialize(const char *filename) { char *gieModelStream{nullptr}; @@ -488,13 +392,89 @@ bool NetworkRT::deserialize(const char *filename) { file.close(); } - PluginFactory plfact; - runtimeRT = createInferRuntime(loggerRT); - engineRT = runtimeRT->deserializeCudaEngine(gieModelStream, size, (IPluginFactory *) &plfact); + engineRT = runtimeRT->deserializeCudaEngine(gieModelStream, size, (IPluginFactory *) pluginFactory); //if (gieModelStream) delete [] gieModelStream; return true; } + + +IPlugin* PluginFactory::createPlugin(const char* layerName, const void* serialData, size_t serialLength) { + const char * buf = reinterpret_cast(serialData); + + std::string name(layerName); + + if(name.find("Activation") == 0) { + 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 + + 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; + } + + 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); + for(int i=0; inum; i++) + r->mask[i] = readBUF(buf); + for(int i=0; i<3*2*r->num; i++) + r->bias[i] = readBUF(buf); + + yolos[n_yolos++] = r; + 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; + } +/* + if(name.find("Route") == 0) { + RouteRT *r = new RouteRT(); + r->in = readBUF(buf); + for(int i=0; ic_in[i] = readBUF(buf); + r->c = readBUF(buf); + r->h = readBUF(buf); + r->w = readBUF(buf); + return r; + } +*/ + FatalError("Cant deserialize Plugin"); + return NULL; +} + }} diff --git a/src/Yolo.cpp b/src/Yolo.cpp index 09e81ce..fc3d3b9 100644 --- a/src/Yolo.cpp +++ b/src/Yolo.cpp @@ -18,13 +18,12 @@ Yolo::Yolo(Network *net, int classes, int num, const char* fname_weights) : this->num = num; // load anchors - int seek = 0; - readBinaryFile(fname_weights, num, &mask_h, &mask_d, seek); - seek += num; - readBinaryFile(fname_weights, 3*num*2, &bias_h, &bias_d, seek); - - printDeviceVector(num, mask_h, false); - printDeviceVector(3*num*2, bias_h, false); + if(fname_weights != nullptr) { + int seek = 0; + readBinaryFile(fname_weights, num, &mask_h, &mask_d, seek); + seek += num; + readBinaryFile(fname_weights, 3*num*2, &bias_h, &bias_d, seek); + } // same output_dim.n = input_dim.n; @@ -33,10 +32,6 @@ Yolo::Yolo(Network *net, int classes, int num, const char* fname_weights) : output_dim.w = input_dim.w; output_dim.l = input_dim.l; - std::cout<<"YOLO INPUT: "; - input_dim.print(); - std::cout<<"\n"; - checkCuda( cudaMalloc(&dstData, output_dim.tot()*sizeof(dnnType)) ); predictions = nullptr; } @@ -62,35 +57,6 @@ Yolo::box get_yolo_box(float *x, float *biases, int n, int index, int i, int j, return b; } -void correct_yolo_boxes(Yolo::detection *dets, int n, int w, int h, int netw, int neth, int relative) -{ - int i; - int new_w=0; - int new_h=0; - if (((float)netw/w) < ((float)neth/h)) { - new_w = netw; - new_h = (h * netw)/w; - } else { - new_h = neth; - new_w = (w * neth)/h; - } - for (i = 0; i < n; ++i){ - Yolo::box b = dets[i].bbox; - b.x = (b.x - (netw - new_w)/2./netw) / ((float)new_w/netw); - b.y = (b.y - (neth - new_h)/2./neth) / ((float)new_h/neth); - b.w *= (float)netw/new_w; - b.h *= (float)neth/new_h; - if(!relative){ - b.x *= w; - b.w *= w; - b.y *= h; - b.h *= h; - } - dets[i].bbox = b; - } -} - - dnnType* Yolo::infer(dataDim_t &dim, dnnType* srcData) { checkCuda( cudaMemcpy(dstData, srcData, dim.tot()*sizeof(dnnType), cudaMemcpyDeviceToDevice)); @@ -109,14 +75,12 @@ dnnType* Yolo::infer(dataDim_t &dim, dnnType* srcData) { return dstData; } -int Yolo::computeDetections(Yolo::detection *dets, int &ndets, int w, int h, int netw, int neth, float thresh) { +int Yolo::computeDetections(Yolo::detection *dets, int &ndets, int netw, int neth, float thresh) { if(predictions == nullptr) predictions = new dnnType[output_dim.tot()]; checkCuda( cudaMemcpy(predictions, dstData, output_dim.tot()*sizeof(dnnType), cudaMemcpyDeviceToHost)); - int relative = 0; - int lw = output_dim.w; int lh = output_dim.h; @@ -150,7 +114,6 @@ int Yolo::computeDetections(Yolo::detection *dets, int &ndets, int w, int h, int } } - correct_yolo_boxes(dets + ndets, count, w, h, netw, neth, relative); ndets = count; return count; } diff --git a/src/Yolo3Detection.cpp b/src/Yolo3Detection.cpp index 5fbdd88..69547b0 100644 --- a/src/Yolo3Detection.cpp +++ b/src/Yolo3Detection.cpp @@ -2,10 +2,36 @@ namespace tk { namespace dnn { -bool Yolo3Detection::init(std::string tensor_folder) { +bool Yolo3Detection::init(std::string tensor_path) { //const char *tensor_path = "../data/yolo3/yolo3_berkeley.rt"; + //convert network to tensorRT + std::cout<<(tensor_path).c_str()<<"\n"; + netRT = new tk::dnn::NetworkRT(NULL, (tensor_path).c_str() ); + + if(netRT->pluginFactory->n_yolos != 3) { + FatalError("this is not yolo3"); + } + + for(int i=0; ipluginFactory->n_yolos; i++) { + YoloRT *yRT = netRT->pluginFactory->yolos[i]; + classes = yRT->classes; + num = yRT->num; + + // make a yolo layer for interpret predictions + yolo[i] = new tk::dnn::Yolo(nullptr, classes, num, nullptr); // yolo without input and bias + memcpy(yolo[i]->mask_h, yRT->mask, sizeof(dnnType)*num); + memcpy(yolo[i]->bias_h, yRT->bias, sizeof(dnnType)*num*3*2); + yolo[i]->input_dim = yolo[i]->output_dim = tk::dnn::dataDim_t(1, yRT->c, yRT->h, yRT->w); + } + + dets = tk::dnn::Yolo::allocateDetections(tk::dnn::Yolo::MAX_DETECTIONS, classes); + + checkCuda(cudaMallocHost(&input, sizeof(dnnType)*netRT->input_dim.tot())); + checkCuda(cudaMalloc(&input_d, sizeof(dnnType)*netRT->input_dim.tot())); + + // class colors precompute for(int c=0; cinput_dim = yolo[0]->output_dim = tk::dnn::dataDim_t(1, 45, 10, 17); - yolo[1] = new tk::dnn::Yolo(nullptr, classes, num, (tensor_folder + "_1.bin").c_str() ); // yolo without input and bias - yolo[1]->input_dim = yolo[1]->output_dim = tk::dnn::dataDim_t(1, 45, 20, 34); - yolo[2] = new tk::dnn::Yolo(nullptr, classes, num, (tensor_folder + "_2.bin").c_str() ); // yolo without input and bias - yolo[2]->input_dim = yolo[2]->output_dim = tk::dnn::dataDim_t(1, 45, 40, 68); - - dets = tk::dnn::Yolo::allocateDetections(tk::dnn::Yolo::MAX_DETECTIONS, classes); - - checkCuda(cudaMallocHost(&input, sizeof(dnnType)*netRT->input_dim.tot())); - checkCuda(cudaMalloc(&input_d, sizeof(dnnType)*netRT->input_dim.tot())); - return true; } @@ -82,7 +91,7 @@ void Yolo3Detection::update(cv::Mat &imageORIG) { for(int i=0; i<3; i++) { rt_out[i] = (dnnType*)netRT->buffersRT[i+1]; yolo[i]->dstData = rt_out[i]; - yolo[i]->computeDetections(dets, ndets, netRT->input_dim.w, netRT->input_dim.h, netRT->input_dim.w, netRT->input_dim.h, thresh); + yolo[i]->computeDetections(dets, ndets, netRT->input_dim.w, netRT->input_dim.h, thresh); } tk::dnn::Yolo::mergeDetections(dets, ndets, classes); TIMER_STOP diff --git a/tests/yolo3_berkeley/yolo3_berkeley.cpp b/tests/yolo3_berkeley/yolo3_berkeley.cpp index 2c947cf..3a69e32 100644 --- a/tests/yolo3_berkeley/yolo3_berkeley.cpp +++ b/tests/yolo3_berkeley/yolo3_berkeley.cpp @@ -326,9 +326,9 @@ int main() { int ndets = 0; int classes = yolo0.classes; tk::dnn::Yolo::detection *dets = tk::dnn::Yolo::allocateDetections(tk::dnn::Yolo::MAX_DETECTIONS, classes); - yolo0.computeDetections(dets, ndets, net.input_dim.w, net.input_dim.h, net.input_dim.w, net.input_dim.h, 0.5); - yolo1.computeDetections(dets, ndets, net.input_dim.w, net.input_dim.h, net.input_dim.w, net.input_dim.h, 0.5); - yolo2.computeDetections(dets, ndets, net.input_dim.w, net.input_dim.h, net.input_dim.w, net.input_dim.h, 0.5); + yolo0.computeDetections(dets, ndets, net.input_dim.w, net.input_dim.h, 0.5); + yolo1.computeDetections(dets, ndets, net.input_dim.w, net.input_dim.h, 0.5); + yolo2.computeDetections(dets, ndets, net.input_dim.w, net.input_dim.h, 0.5); tk::dnn::Yolo::mergeDetections(dets, ndets, classes); for(int j=0; j