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

Additional support for custom network architectures #176

Open
nightduck wants to merge 3 commits from nightduck/custom_arch_support into master
pull from: nightduck/custom_arch_support
merge into: mmr:master
Conversation 3 Commits 3 Files Changed 3
+44 -17
3 changed files with 44 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files
include/tkDNN/Network.h
+33
View File
@@ -3,9 +3,16 @@
#include <string>
#include "utils.h"
#include "NvInfer.h"
namespace tk { namespace dnn {
enum dimFormat_t {
CHW,
NCHW,
//NHWC
};
/**
Data representation between layers
n = batch size
@@ -20,6 +27,32 @@ struct dataDim_t {
dataDim_t() : n(1), c(1), h(1), w(1), l(1) {};
dataDim_t(nvinfer1::Dims &d, dimFormat_t df) {
switch(df) {
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) :
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;
NetworkRT(Network *net, const char *name);
NetworkRT(Network *net, const char *name, dimFormat_t dim_format=CHW, const char *input_name="data", const char *output_name="out");
virtual ~NetworkRT();
int getMaxBatchSize() {
src/NetworkRT.cpp
+10 -16
View File
@@ -26,7 +26,7 @@ namespace tk { namespace dnn {
std::map<Layer*, nvinfer1::ITensor*>tensors;
NetworkRT::NetworkRT(Network *net, const char *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(NV_TENSORRT_MINOR)/10 +
@@ -97,13 +97,13 @@ NetworkRT::NetworkRT(Network *net, const char *name) {
calibrator.reset(new Int8EntropyCalibrator(calibrationStream, 1,
calib_table_name,
"data"));
input_name));
configRT->setInt8Calibrator(calibrator.get());
}
#endif
// add input layer
ITensor *input = networkRT->addInput("data", DataType::kFLOAT,
ITensor *input = networkRT->addInput(input_name, DataType::kFLOAT,
DimsCHW{ dim.c, dim.h, dim.w});
checkNULL(input);
@@ -130,7 +130,7 @@ NetworkRT::NetworkRT(Network *net, const char *name) {
FatalError("conversion failed");
//build tensorRT
input->setName("out");
input->setName(output_name);
networkRT->markOutput(*input);
std::cout<<"Selected maxBatchSize: "<<builderRT->getMaxBatchSize()<<"\n";
@@ -161,31 +161,25 @@ NetworkRT::NetworkRT(Network *net, const char *name) {
// In order to bind the buffers, we need to know the names of the input and output tensors.
// note that indices are guaranteed to be less than IEngine::getNbBindings()
buf_input_idx = engineRT->getBindingIndex("data");
buf_output_idx = engineRT->getBindingIndex("out");
buf_input_idx = engineRT->getBindingIndex(input_name);
buf_output_idx = engineRT->getBindingIndex(output_name);
std::cout<<"input index = "<<buf_input_idx<<" -> output index = "<<buf_output_idx<<"\n";
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, dim_format);
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, dim_format);
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, dim_format);
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));