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

F16 inference

Browse Source
This commit is contained in:
Francesco Gatti
2017-08-14 10:16:29 +00:00
parent 66ad6bb1d6
commit 2b4b9b8e49
1 changed files with 20 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/NetworkRT.cpp
+20 -12
View File
@@ -36,12 +36,29 @@ NetworkRT::NetworkRT(Network *net, const char *name) {
std::cout<<"Float16 support: "<<builderRT->platformHasFastFp16()<<"\n";
std::cout<<"Int8 support: "<<builderRT->platformHasFastInt8()<<"\n";
networkRT = builderRT->createNetwork();
dtRT = DataType::kFLOAT;
if(!fileExist(name)) {
//add input layer
//input and dataType
dataDim_t dim = net->layers[0]->input_dim;
dtRT = DataType::kFLOAT;
builderRT->setMaxBatchSize(1);
builderRT->setMaxWorkspaceSize(1 << 30);
//change datatype based on system specs
if(builderRT->platformHasFastInt8()) {
BatchStream bstream({32,dim.c, dim.h, dim.w}, 32, 1);
Int8EntropyCalibrator calib(bstream, 0, false);
builderRT->setInt8Mode(true);
builderRT->setInt8Calibrator(&calib);
} else if(builderRT->platformHasFastFp16()) {
dtRT = DataType::kHALF;
builderRT->setHalf2Mode(true);
}
//add input layer
ITensor *input = networkRT->addInput("data", dtRT,
DimsCHW{ dim.c, dim.h, dim.w});
checkNULL(input);
@@ -62,15 +79,6 @@ NetworkRT::NetworkRT(Network *net, const char *name) {
input->setName("out");
networkRT->markOutput(*input);
// 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