Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Additional support for custom network architectures #176

Open
wants to merge 3 commits into
base: master
Choose a base branch
from
Open

Additional support for custom network architectures #176

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
714d068
Initiliaze dataDim_t with nvinfer1::Dims
nightduck Dec 5, 2020
f6884f0
Can now specify input and output names to NetworkRT
nightduck Dec 15, 2020
04a5823
Added options to specify dimension ordering
nightduck Jan 12, 2021
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
33 changes: 33 additions & 0 deletions include/tkDNN/Network.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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) {};

Expand Down
2 changes: 1 addition & 1 deletion include/tkDNN/NetworkRT.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -73,7 +73,7 @@ class NetworkRT {

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() {
Expand Down
26 changes: 10 additions & 16 deletions src/NetworkRT.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 +
Expand Down Expand Up @@ -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);

Expand All @@ -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";
Expand Down Expand Up @@ -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));
Expand Down