mmr/tkDNN
Archived
1
0
Fork 0
Code Issues 46 Pull Requests 12 Releases 4 Activity

Add GPU version for yolo3detection

Browse Source 
Signed-off-by: Micaela Verucchi <micaelaverucchi@gmail.com>
This commit is contained in:
Micaela Verucchi
2020-03-23 21:29:29 +01:00
parent df37e11709
commit 3eda9b9219
3 changed files with 11 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
include/tkDNN/DetectionNN.h
+2 -2
View File
@@ -3,7 +3,7 @@
#include <iostream> #include <iostream>
#include <signal.h> #include <signal.h>
#include <stdlib.h> /* srand, rand */ #include <stdlib.h>
#include <unistd.h> #include <unistd.h>
#include <mutex> #include <mutex>
#include "utils.h" #include "utils.h"
@@ -14,7 +14,7 @@
#include "tkdnn.h" #include "tkdnn.h"
// #define OPENCV_CUDA //if OPENCV has been compiled with CUDA and contrib. #define OPENCV_CUDA //if OPENCV has been compiled with CUDA and contrib.
namespace tk { namespace dnn { namespace tk { namespace dnn {
include/tkDNN/Yolo3Detection.h
+2
View File
@@ -17,6 +17,8 @@ private:
tk::dnn::Yolo* yolo[3]; tk::dnn::Yolo* yolo[3];
tk::dnn::Yolo* getYoloLayer(int n=0); tk::dnn::Yolo* getYoloLayer(int n=0);
cv::Mat bgr_h;
public: public:
Yolo3Detection() {}; Yolo3Detection() {};
src/Yolo3Detection.cpp
+7 -3
View File
@@ -62,9 +62,10 @@ void Yolo3Detection::preprocess(cv::Mat &frame)
//write channels //write channels
for(int i=0; i<netRT->input_dim.c; i++) { for(int i=0; i<netRT->input_dim.c; i++) {
int idx = i*imagePreproc.rows*imagePreproc.cols; int size = imagePreproc.rows * imagePreproc.cols;
int ch = netRT->input_dim.c-1 -i; int ch = netRT->input_dim.c-1 -i;
checkCuda( cudaMemcpy((void*)&input_d[idx], (void*)bgr[ch].data, imagePreproc.rows*imagePreproc.cols*sizeof(dnnType), cudaMemcpyDeviceToDevice)); bgr[ch].download(bgr_h); //TODO: don't copy back on CPU
checkCuda( cudaMemcpy(input_d + i*size, (float*)bgr_h.data, size*sizeof(dnnType), cudaMemcpyHostToDevice));
} }
#else #else
cv::resize(frame, frame, cv::Size(netRT->input_dim.w, netRT->input_dim.h)); cv::resize(frame, frame, cv::Size(netRT->input_dim.w, netRT->input_dim.h));
@@ -77,7 +78,7 @@ void Yolo3Detection::preprocess(cv::Mat &frame)
for(int i=0; i<netRT->input_dim.c; i++) { for(int i=0; i<netRT->input_dim.c; i++) {
int idx = i*imagePreproc.rows*imagePreproc.cols; int idx = i*imagePreproc.rows*imagePreproc.cols;
int ch = netRT->input_dim.c-1 -i; int ch = netRT->input_dim.c-1 -i;
memcpy((void*)&input[idx], (void*)bgr[ch].data, imagePreproc.rows*imagePreproc.cols*sizeof(dnnType)); memcpy((void*)&input[idx], (void*)bgr[ch].data, imagePreproc.rows*imagePreproc.cols*sizeof(dnnType));
} }
checkCuda(cudaMemcpyAsync(input_d, input, netRT->input_dim.tot()*sizeof(dnnType), cudaMemcpyHostToDevice, netRT->stream)); checkCuda(cudaMemcpyAsync(input_d, input, netRT->input_dim.tot()*sizeof(dnnType), cudaMemcpyHostToDevice, netRT->stream));
#endif #endif
@@ -96,6 +97,9 @@ void Yolo3Detection::update(cv::Mat &frame)
//do inference //do inference
tk::dnn::dataDim_t dim = netRT->input_dim; tk::dnn::dataDim_t dim = netRT->input_dim;
// printDeviceVector(netRT->input_dim.tot()*sizeof(dnnType),input_d);
printCenteredTitle(" TENSORRT inference ", '=', 30); printCenteredTitle(" TENSORRT inference ", '=', 30);
{ {