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

Added options to specify dimension ordering

Browse Source
This commit is contained in:
oren
2021-01-12 09:43:54 -06:00
parent f6884f07bf
commit 04a5823873
3 changed files with 36 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files
include/tkDNN/Network.h
+31 -3
View File
@@ -7,6 +7,12 @@
namespace tk { namespace dnn { namespace tk { namespace dnn {
enum dimFormat_t {
CHW,
NCHW,
//NHWC
};
/** /**
Data representation between layers Data representation between layers
n = batch size n = batch size
@@ -21,9 +27,31 @@ struct dataDim_t {
dataDim_t() : n(1), c(1), h(1), w(1), l(1) {}; dataDim_t() : n(1), c(1), h(1), w(1), l(1) {};
dataDim_t(nvinfer1::Dims &d) : dataDim_t(nvinfer1::Dims &d, dimFormat_t df) {
n(1), c(d.d[0] ? d.d[0] : 1), h(d.d[1] ? d.d[1] : 1), switch(df) {
w(d.d[2] ? d.d[2] : 1), l(d.d[3] ? d.d[3] : 1) {}; case CHW:
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;
break;
case NCHW:
n = d.d[0] ? d.d[0] : 1;
c = d.d[1] ? d.d[1] : 1;
h = d.d[2] ? d.d[2] : 1;
w = d.d[3] ? d.d[3] : 1;
l = d.d[4] ? d.d[4] : 1;
break;
// case NHWC:
// n = d.d[0] ? d.d[0] : 1;
// h = d.d[1] ? d.d[1] : 1;
// w = d.d[2] ? d.d[2] : 1;
// c = d.d[3] ? d.d[3] : 1;
// l = d.d[4] ? d.d[4] : 1;
// break;
}
};
dataDim_t(int _n, int _c, int _h, int _w, int _l = 1) : dataDim_t(int _n, int _c, int _h, int _w, int _l = 1) :
n(_n), c(_c), h(_h), w(_w), l(_l) {}; n(_n), c(_c), h(_h), w(_w), l(_l) {};
include/tkDNN/NetworkRT.h
+1 -1
View File
@@ -73,7 +73,7 @@ public:
PluginFactory *pluginFactory; PluginFactory *pluginFactory;
NetworkRT(Network *net, const char *name, const char *input_name="data", const char *output_name="out"); NetworkRT(Network *net, const char *name, dimFormat_t dim_format=CHW, const char *input_name="data", const char *output_name="out");
virtual ~NetworkRT(); virtual ~NetworkRT();
int getMaxBatchSize() { int getMaxBatchSize() {
src/NetworkRT.cpp
+4 -4
View File
@@ -26,7 +26,7 @@ namespace tk { namespace dnn {
std::map<Layer*, nvinfer1::ITensor*>tensors; std::map<Layer*, nvinfer1::ITensor*>tensors;
NetworkRT::NetworkRT(Network *net, const char *name, const char *input_name, const char *output_name) { NetworkRT::NetworkRT(Network *net, const char *name, dimFormat_t dim_format, const char *input_name, const char *output_name) {
float rt_ver = float(NV_TENSORRT_MAJOR) + float rt_ver = float(NV_TENSORRT_MAJOR) +
float(NV_TENSORRT_MINOR)/10 + float(NV_TENSORRT_MINOR)/10 +
@@ -167,17 +167,17 @@ NetworkRT::NetworkRT(Network *net, const char *name, const char *input_name, con
Dims iDim = engineRT->getBindingDimensions(buf_input_idx); Dims iDim = engineRT->getBindingDimensions(buf_input_idx);
input_dim = dataDim_t(iDim); input_dim = dataDim_t(iDim, dim_format);
input_dim.print(); input_dim.print();
Dims oDim = engineRT->getBindingDimensions(buf_output_idx); Dims oDim = engineRT->getBindingDimensions(buf_output_idx);
output_dim = dataDim_t(oDim); output_dim = dataDim_t(oDim, dim_format);
output_dim.print(); output_dim.print();
// create GPU buffers and a stream // create GPU buffers and a stream
for(int i=0; i<engineRT->getNbBindings(); i++) { for(int i=0; i<engineRT->getNbBindings(); i++) {
Dims dim = engineRT->getBindingDimensions(i); Dims dim = engineRT->getBindingDimensions(i);
buffersDIM[i] = dataDim_t(dim); buffersDIM[i] = dataDim_t(dim, dim_format);
std::cout<<"RtBuffer "<<i<<" dim: "; buffersDIM[i].print(); std::cout<<"RtBuffer "<<i<<" dim: "; buffersDIM[i].print();
checkCuda(cudaMalloc(&buffersRT[i], engineRT->getMaxBatchSize()*buffersDIM[i].tot()*sizeof(dnnType))); checkCuda(cudaMalloc(&buffersRT[i], engineRT->getMaxBatchSize()*buffersDIM[i].tot()*sizeof(dnnType)));
} }