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

input dim fix

Browse Source
This commit is contained in:
Francesco Gatti
2017-08-14 11:57:28 +02:00
parent fc9fb4f153
commit 66ad6bb1d6
1 changed files with 4 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/NetworkRT.cpp
+4 -3
View File
@@ -65,12 +65,12 @@ NetworkRT::NetworkRT(Network *net, const char *name) {
// Build the engine
builderRT->setMaxBatchSize(1);
builderRT->setMaxWorkspaceSize(1 << 20);
/*
BatchStream bstream({32,dim.c, dim.h, dim.w}, 32, 1);
Int8EntropyCalibrator calib(bstream, 0, false);
builderRT->setInt8Mode(true);
builderRT->setInt8Calibrator(&calib);
*/
std::cout<<"Building tensorRT cuda engine...\n";
engineRT = builderRT->buildCudaEngine(*networkRT);
// we don't need the network any more
@@ -95,11 +95,12 @@ NetworkRT::NetworkRT(Network *net, const char *name) {
std::cout<<"input idex = "<<buf_input_idx<<" -> output index = "<<buf_output_idx<<"\n";
Dims iDim = engineRT->getBindingDimensions(buf_output_idx);
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.print();
Dims oDim = engineRT->getBindingDimensions(buf_output_idx);
output_dim.n = 1;