diff --git a/demo/demo/demo.cpp b/demo/demo/demo.cpp index ac6c63d..295b2af 100644 --- a/demo/demo/demo.cpp +++ b/demo/demo/demo.cpp @@ -25,7 +25,7 @@ int main(int argc, char *argv[]) { signal(SIGINT, sig_handler); - char *net = "yolo3_berkeley.rt"; + char *net = "yolo3.rt"; if(argc > 1) net = argv[1]; char *input = "../demo/yolo_test.mp4"; diff --git a/include/tkDNN/NetworkRT.h b/include/tkDNN/NetworkRT.h index d30da03..5cdc34c 100644 --- a/include/tkDNN/NetworkRT.h +++ b/include/tkDNN/NetworkRT.h @@ -30,6 +30,7 @@ using namespace nvinfer1; #include "pluginsRT/ShortcutRT.h" #include "pluginsRT/YoloRT.h" #include "pluginsRT/UpsampleRT.h" +#include "pluginsRT/ResizeLayerRT.h" //#include "pluginsRT/Int8Calibrator.h" class PluginFactory : IPluginFactory diff --git a/include/tkDNN/Yolo3Detection.h b/include/tkDNN/Yolo3Detection.h index 0a7ba25..6c873bc 100644 --- a/include/tkDNN/Yolo3Detection.h +++ b/include/tkDNN/Yolo3Detection.h @@ -33,6 +33,7 @@ class Yolo3Detection { public: int classes = 0; int num = 0; + int n_masks = 0; float thresh = 0.3; cv::Scalar colors[256]; diff --git a/include/tkDNN/kernels.h b/include/tkDNN/kernels.h index dfff6e3..42e464f 100644 --- a/include/tkDNN/kernels.h +++ b/include/tkDNN/kernels.h @@ -9,6 +9,9 @@ void activationLOGISTICForward(dnnType* srcData, dnnType* dstData, int size, cud void fill(dnnType* data, int size, dnnType val, cudaStream_t stream = cudaStream_t(0)); +void resizeForward( dnnType* srcData, dnnType* dstData, int n, int i_c, int i_h, int i_w, + int o_c, int o_h, int o_w, cudaStream_t stream = cudaStream_t(0)); + void reorgForward( dnnType* srcData, dnnType* dstData, int n, int c, int h, int w, int stride, cudaStream_t stream = cudaStream_t(0)); void softmaxForward(float *input, int n, int batch, int batch_offset, diff --git a/include/tkDNN/pluginsRT/ResizeLayerRT.h b/include/tkDNN/pluginsRT/ResizeLayerRT.h new file mode 100644 index 0000000..ae87dbf --- /dev/null +++ b/include/tkDNN/pluginsRT/ResizeLayerRT.h @@ -0,0 +1,67 @@ +#include +#include "../kernels.h" + +class ResizeLayerRT : public IPlugin { + +public: + ResizeLayerRT(int c, int h, int w) { + o_c = c; + o_h = h; + o_w = w; + } + + ~ResizeLayerRT(){ + } + + int getNbOutputs() const override { + return 1; + } + + Dims getOutputDimensions(int index, const Dims* inputs, int nbInputDims) override { + return DimsCHW{o_c, o_h, o_w}; + } + + void configure(const Dims* inputDims, int nbInputs, const Dims* outputDims, int nbOutputs, int maxBatchSize) override { + i_c = inputDims[0].d[0]; + i_h = inputDims[0].d[1]; + i_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 { + // printf("%d %d %d %d %d %d\n", i_c, i_w, i_h, o_c, o_w, o_h); + resizeForward((dnnType*)reinterpret_cast(inputs[0]), + reinterpret_cast(outputs[0]), + batchSize, i_c, i_h, i_w, o_c, o_h, o_w, stream); + return 0; + } + + + virtual size_t getSerializationSize() override { + return 6*sizeof(int); + } + + virtual void serialize(void* buffer) override { + char *buf = reinterpret_cast(buffer); + + tk::dnn::writeBUF(buf, o_c); + tk::dnn::writeBUF(buf, o_h); + tk::dnn::writeBUF(buf, o_w); + + tk::dnn::writeBUF(buf, i_c); + tk::dnn::writeBUF(buf, i_h); + tk::dnn::writeBUF(buf, i_w); + } + + int i_c, i_h, i_w, o_c, o_h, o_w; +}; diff --git a/src/NetworkRT.cpp b/src/NetworkRT.cpp index 4500be5..e31185a 100644 --- a/src/NetworkRT.cpp +++ b/src/NetworkRT.cpp @@ -15,9 +15,9 @@ using namespace nvinfer1; // Logger for info/warning/errors class Logger : public ILogger { void log(Severity severity, const char* msg) override { -// #ifdef DEBUG +#ifdef DEBUG std::cout <<"TENSORRT LOG: "<< msg << std::endl; -// #endif +#endif } } loggerRT; @@ -209,8 +209,8 @@ ILayer* NetworkRT::convert_layer(ITensor *input, Dense *l) { ILayer* NetworkRT::convert_layer(ITensor *input, Conv2d *l) { - std::cout<<"convert conv2D\n"; - printf("%d %d %d %d %d\n", l->kernelH, l->kernelW, l->inputs, l->outputs, l->batchnorm); + // std::cout<<"convert conv2D\n"; + // printf("%d %d %d %d %d\n", l->kernelH, l->kernelW, l->inputs, l->outputs, l->batchnorm); void *data_b, *bias_b, *power_b, *mean_b, *variance_b, *scales_b; @@ -261,7 +261,7 @@ ILayer* NetworkRT::convert_layer(ITensor *input, Conv2d *l) { Weights power{dtRT, power_b, l->outputs}; Weights shift{dtRT, mean_b, l->outputs}; Weights scale{dtRT, variance_b, l->outputs}; - std::cout<getNbOutputs()<getNbOutputs()<addScale(*lRT->getOutput(0), ScaleMode::kCHANNEL, shift, scale, power); @@ -280,8 +280,7 @@ ILayer* NetworkRT::convert_layer(ITensor *input, Conv2d *l) { } ILayer* NetworkRT::convert_layer(ITensor *input, Pooling *l) { - std::cout<<"convert Pooling\n"; - printf("%d %d %d %d %d %d %d %d %d %d %d %d (layer)\n", l->input_dim.h, l->input_dim.w, l->output_dim.h, l->output_dim.w, l->winW, l->winH, l->strideH, l->strideW, l->paddingH, l->paddingW, l->pool_mode, tkdnnPoolingMode_t::POOLING_MAX) ; + // std::cout<<"convert Pooling\n"; PoolingType ptype; if(l->pool_mode == tkdnnPoolingMode_t::POOLING_MAX) ptype = PoolingType::kMAX; @@ -291,29 +290,28 @@ ILayer* NetworkRT::convert_layer(ITensor *input, Pooling *l) { IPoolingLayer *lRT = networkRT->addPooling(*input, ptype, DimsHW{l->winH, l->winW}); checkNULL(lRT); + + lRT->setPadding(DimsHW{l->paddingH, l->paddingW}); + lRT->setStride(DimsHW{l->strideH, l->strideW}); + - // if (l->input_dim.h == 13 && l->output_dim.h == 13) - // { - // lRT->setPadding(DimsHW{7, 7}); - // lRT->setStride(DimsHW{2, 2}); - // } - // else - // { - lRT->setPadding(DimsHW{l->paddingH, l->paddingW}); - lRT->setStride(DimsHW{l->strideH, l->strideW}); - // } - - // IResizeLayer *lRT = networkRT->addResize(*lRT->getOutput(0)); - // checkNULL(lRT); - // lRT->setOutputDimensions(l->output_dim); - ITensor *t = lRT->getOutput(0); - for(int j=0; jgetDimensions().nbDims; j++) { - std::cout<getDimensions().d[j]<<" "; - } - std::cout<<" (TensorRT)\n"; + // for(int j=0; jgetDimensions().nbDims; j++) { + // std::cout<getDimensions().d[j]<<" "; + // } + // std::cout<<" (TensorRT)\n"; - return lRT; + IPlugin *plugin = new ResizeLayerRT( l->output_dim.c,l->output_dim.h,l->output_dim.w ); + IPluginLayer *lRT1 = networkRT->addPlugin(&t, 1, *plugin); + checkNULL(lRT1); + + // ITensor *t1 = lRT1->getOutput(0); + // for(int j=0; jgetDimensions().nbDims; j++) { + // std::cout<getDimensions().d[j]<<" "; + // } + // std::cout<<" (TensorRT after resize )\n"; + + return lRT1; } ILayer* NetworkRT::convert_layer(ITensor *input, Activation *l) { @@ -347,17 +345,17 @@ ILayer* NetworkRT::convert_layer(ITensor *input, Softmax *l) { } ILayer* NetworkRT::convert_layer(ITensor *input, Route *l) { - std::cout<<"convert route\n"; + // std::cout<<"convert route\n"; ITensor **tens = new ITensor*[l->layers_n]; for(int i=0; ilayers_n; i++) { tens[i] = tensors[l->layers[i]]; - for(int j=0; jgetDimensions().nbDims; j++) { - std::cout<getDimensions().d[j]<<" "; - } - std::cout<<"\n"; + // for(int j=0; jgetDimensions().nbDims; j++) { + // std::cout<getDimensions().d[j]<<" "; + // } + // std::cout<<"\n"; } IConcatenationLayer *lRT = networkRT->addConcatenation(tens, l->layers_n); @@ -502,6 +500,16 @@ IPlugin* PluginFactory::createPlugin(const char* layerName, const void* serialDa return r; } + if(name.find("Pooling") == 0) { + ResizeLayerRT *r = new ResizeLayerRT(readBUF(buf), //o_c + readBUF(buf), //o_h + readBUF(buf)); //o_w + r->i_c = readBUF(buf); + r->i_h = readBUF(buf); + r->i_w = readBUF(buf); + return r; + } + if(name.find("Yolo") == 0) { YoloRT *r = new YoloRT(readBUF(buf), //classes readBUF(buf), //num diff --git a/src/Pooling.cpp b/src/Pooling.cpp index ce84926..0f3573d 100644 --- a/src/Pooling.cpp +++ b/src/Pooling.cpp @@ -53,7 +53,6 @@ Pooling::Pooling( Network *net, int winH, int winW, int strideH, int strideW, // checkCUDNN( cudnnGetPooling2dForwardOutputDim(poolingDesc, srcTensorDesc, &n, &c, &h, &w)); //compute w and h as in darknet - int padH = paddingH == 0? winH -1 : paddingH; int padW = paddingW == 0? winW -1 : paddingW; diff --git a/src/Yolo.cpp b/src/Yolo.cpp index ae4694c..ca6ebc5 100644 --- a/src/Yolo.cpp +++ b/src/Yolo.cpp @@ -130,7 +130,7 @@ int Yolo::computeDetections(Yolo::detection *dets, int &ndets, int netw, int net for (i = 0; i < lw*lh; ++i){ int row = i / lw; int col = i % lw; - for(n = 0; n < num; ++n){ + for(n = 0; n < n_masks; ++n){ int obj_index = entry_index(0, n*lw*lh + i, 4, classes, input_dim, output_dim); float objectness = predictions[obj_index]; if(objectness <= thresh) continue; diff --git a/src/Yolo3Detection.cpp b/src/Yolo3Detection.cpp index ed392f7..aae23b7 100644 --- a/src/Yolo3Detection.cpp +++ b/src/Yolo3Detection.cpp @@ -22,7 +22,8 @@ bool Yolo3Detection::init(std::string tensor_path) { std::cout<<(tensor_path).c_str()<<"\n"; netRT = new tk::dnn::NetworkRT(NULL, (tensor_path).c_str() ); - if(netRT->pluginFactory->n_yolos != 3) { + + if(netRT->pluginFactory->n_yolos < 2 ) { FatalError("this is not yolo3"); } @@ -30,13 +31,14 @@ bool Yolo3Detection::init(std::string tensor_path) { YoloRT *yRT = netRT->pluginFactory->yolos[i]; classes = yRT->classes; num = yRT->num; + n_masks = yRT->n_masks; // make a yolo layer for interpret predictions - yolo[i] = new tk::dnn::Yolo(nullptr, classes, num, ""); // yolo without input and bias - yolo[i]->mask_h = new dnnType[num]; - yolo[i]->bias_h = new dnnType[num*3*2]; - memcpy(yolo[i]->mask_h, yRT->mask, sizeof(dnnType)*num); - memcpy(yolo[i]->bias_h, yRT->bias, sizeof(dnnType)*num*3*2); + yolo[i] = new tk::dnn::Yolo(nullptr, classes, n_masks, ""); // yolo without input and bias + yolo[i]->mask_h = new dnnType[n_masks]; + yolo[i]->bias_h = new dnnType[num*n_masks*2]; + memcpy(yolo[i]->mask_h, yRT->mask, sizeof(dnnType)*n_masks); + memcpy(yolo[i]->bias_h, yRT->bias, sizeof(dnnType)*num*n_masks*2); yolo[i]->input_dim = yolo[i]->output_dim = tk::dnn::dataDim_t(1, yRT->c, yRT->h, yRT->w); yolo[i]->classesNames = yRT->classesNames; } @@ -84,7 +86,7 @@ void Yolo3Detection::update(cv::Mat &imageORIG) { //DO INFERENCE - dnnType *rt_out[3]; + dnnType *rt_out[netRT->pluginFactory->n_yolos]; tk::dnn::dataDim_t dim = netRT->input_dim; checkCuda(cudaMemcpyAsync(input_d, input, dim.tot()*sizeof(dnnType), cudaMemcpyHostToDevice, netRT->stream)); @@ -99,7 +101,7 @@ void Yolo3Detection::update(cv::Mat &imageORIG) { TIMER_START // compute dets ndets = 0; - for(int i=0; i<3; i++) { + for(int i=0; ipluginFactory->n_yolos; 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, thresh); diff --git a/src/kernels/resize.cu b/src/kernels/resize.cu new file mode 100644 index 0000000..f637152 --- /dev/null +++ b/src/kernels/resize.cu @@ -0,0 +1,46 @@ +#include "kernels.h" +#include +#define MIN(a,b) (((a)<(b))?(a):(b)) +#define MAX(a,b) (((a)>(b))?(a):(b)) + +__global__ void resize_kernel( int i_N,float *x, int i_w, int i_h, int i_c, + int o_w, int o_h, int o_c, int batch, float *out) +{ + int i = (blockIdx.x + blockIdx.y*gridDim.x) * blockDim.x + threadIdx.x; + if(i >= i_N) return; + + int out_index = i; + int out_w = i%o_w; + i = i/o_w; + int out_h = i%o_h; + i = i/o_h; + int out_c = i%o_c; + i = i/o_c; + + //copying last column/last row + int in_index = ((i*i_c + MIN(out_c,i_c-1))*i_h + MIN(out_h,i_h-1))*i_w + MIN(out_w, i_w-1); + out[out_index] = x[in_index]; +} + + +void resizeForward( dnnType* srcData, dnnType* dstData, int n, int i_c, int i_h, int i_w, + int o_c, int o_h, int o_w, cudaStream_t stream ) +{ + int i_size = n*i_c*i_h*i_w; + int o_size = n*o_c*o_h*o_w; + + int blocks = (o_size+255)/256; + int threads = 256; + + if(i_c == o_c && i_h == o_h && i_w == o_w ) + { + checkCuda(cudaMemcpy(dstData, srcData, i_size*sizeof(dnnType), cudaMemcpyDeviceToDevice)); + } + else + { + checkCuda(cudaMemset(dstData, 0, o_size*sizeof(dnnType))); + resize_kernel<<>>(o_size, srcData, i_w, i_h, i_c, o_w, o_h, o_c, n, dstData); + // printDeviceVector(i_size, srcData); + // printDeviceVector(o_size, dstData); + } +} \ No newline at end of file diff --git a/tests/yolo3_tiny/yolo3_tiny.cpp b/tests/yolo3_tiny/yolo3_tiny.cpp index c04b37d..7c8a382 100644 --- a/tests/yolo3_tiny/yolo3_tiny.cpp +++ b/tests/yolo3_tiny/yolo3_tiny.cpp @@ -118,5 +118,6 @@ int main() { 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; }