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

Initiliaze dataDim_t with nvinfer1::Dims

Browse Source
This commit is contained in:
oren
2020-12-05 16:52:47 -06:00
parent 86478f9384
commit 714d0685a9
2 changed files with 9 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files
include/tkDNN/Network.h
+5
View File
@@ -3,6 +3,7 @@
#include <string>
#include "utils.h"
#include "NvInfer.h"
namespace tk { namespace dnn {
@@ -20,6 +21,10 @@ struct dataDim_t {
dataDim_t() : n(1), c(1), h(1), w(1), l(1) {};
dataDim_t(nvinfer1::Dims &d) :
n(1), c(d.d[0] ? d.d[0] : 1), h(d.d[1] ? d.d[1] : 1),
w(d.d[2] ? d.d[2] : 1), l(d.d[3] ? d.d[3] : 1) {};
dataDim_t(int _n, int _c, int _h, int _w, int _l = 1) :
n(_n), c(_c), h(_h), w(_w), l(_l) {};
src/NetworkRT.cpp
+4 -10
View File
@@ -167,25 +167,19 @@ NetworkRT::NetworkRT(Network *net, const char *name) {
Dims iDim = engineRT->getBindingDimensions(buf_input_idx);
input_dim.n = 1;
input_dim.c = iDim.d[0];
input_dim.h = iDim.d[1];
input_dim.w = iDim.d[2];
input_dim = dataDim_t(iDim);
input_dim.print();
Dims oDim = engineRT->getBindingDimensions(buf_output_idx);
output_dim.n = 1;
output_dim.c = oDim.d[0];
output_dim.h = oDim.d[1];
output_dim.w = oDim.d[2];
output_dim = dataDim_t(oDim);
output_dim.print();
// create GPU buffers and a stream
for(int i=0; i<engineRT->getNbBindings(); i++) {
Dims dim = engineRT->getBindingDimensions(i);
buffersDIM[i] = dataDim_t(1, dim.d[0], dim.d[1], dim.d[2]);
buffersDIM[i] = dataDim_t(dim);
std::cout<<"RtBuffer "<<i<<" dim: "; buffersDIM[i].print();
checkCuda(cudaMalloc(&buffersRT[i], engineRT->getMaxBatchSize()*dim.d[0]*dim.d[1]*dim.d[2]*sizeof(dnnType)));
checkCuda(cudaMalloc(&buffersRT[i], engineRT->getMaxBatchSize()*buffersDIM[i].tot()*sizeof(dnnType)));
}
checkCuda(cudaMalloc(&output, engineRT->getMaxBatchSize()*output_dim.tot()*sizeof(dnnType)));
checkCuda(cudaStreamCreate(&stream));