diff --git a/include/tkDNN/NetworkRT.h b/include/tkDNN/NetworkRT.h index abae6e5..64deac9 100644 --- a/include/tkDNN/NetworkRT.h +++ b/include/tkDNN/NetworkRT.h @@ -55,8 +55,9 @@ public: nvinfer1::IBuilder *builderRT; nvinfer1::IRuntime *runtimeRT; nvinfer1::INetworkDefinition *networkRT; +#if NV_TENSORRT_MAJOR >= 6 nvinfer1::IBuilderConfig *configRT; - +#endif nvinfer1::ICudaEngine *engineRT; nvinfer1::IExecutionContext *contextRT; diff --git a/src/NetworkRT.cpp b/src/NetworkRT.cpp index a22b3eb..bf74010 100644 --- a/src/NetworkRT.cpp +++ b/src/NetworkRT.cpp @@ -38,9 +38,12 @@ NetworkRT::NetworkRT(Network *net, const char *name) { std::cout<<"Int8 support: "<platformHasFastInt8()<<"\n"; std::cout<<"DLAs: "<getNbDLACores()<<"\n"; networkRT = builderRT->createNetwork(); - configRT = builderRT->createBuilderConfig(); +#if NV_TENSORRT_MAJOR >= 6 + configRT = builderRT->createBuilderConfig(); +#endif if(!fileExist(name)) { +#if NV_TENSORRT_MAJOR >= 6 // Calibrator life time needs to last until after the engine is built. std::unique_ptr calibrator; @@ -48,7 +51,7 @@ NetworkRT::NetworkRT(Network *net, const char *name) { configRT->setMinTimingIterations(1); configRT->setMaxWorkspaceSize(1 << 30); configRT->setFlag(BuilderFlag::kDEBUG); - +#endif //input and dataType dataDim_t dim = net->layers[0]->input_dim; dtRT = DataType::kFLOAT; @@ -59,7 +62,9 @@ NetworkRT::NetworkRT(Network *net, const char *name) { if(net->fp16 && builderRT->platformHasFastFp16()) { dtRT = DataType::kHALF; builderRT->setHalf2Mode(true); +#if NV_TENSORRT_MAJOR >= 6 configRT->setFlag(BuilderFlag::kFP16); +#endif } if(net->dla && builderRT->getNbDLACores() > 0) { dtRT = DataType::kHALF; @@ -68,6 +73,7 @@ NetworkRT::NetworkRT(Network *net, const char *name) { builderRT->setDefaultDeviceType(DeviceType::kDLA); builderRT->setDLACore(0); } +#if NV_TENSORRT_MAJOR >= 6 if(net->int8 && builderRT->platformHasFastInt8()){ // dtRT = DataType::kINT8; // builderRT->setInt8Mode(true); @@ -90,6 +96,7 @@ NetworkRT::NetworkRT(Network *net, const char *name) { "data")); configRT->setInt8Calibrator(calibrator.get()); } +#endif // add input layer ITensor *input = networkRT->addInput("data", DataType::kFLOAT, @@ -100,10 +107,12 @@ NetworkRT::NetworkRT(Network *net, const char *name) { for(int i=0; inum_layers; i++) { Layer *l = net->layers[i]; ILayer *Ilay = convert_layer(input, l); +#if NV_TENSORRT_MAJOR >= 6 if(net->int8 && builderRT->platformHasFastInt8()) { Ilay->setPrecision(DataType::kINT8); } +#endif Ilay->setName( (l->getLayerName() + std::to_string(i)).c_str() ); input = Ilay->getOutput(0); @@ -121,7 +130,11 @@ NetworkRT::NetworkRT(Network *net, const char *name) { networkRT->markOutput(*input); std::cout<<"Building tensorRT cuda engine...\n"; +#if NV_TENSORRT_MAJOR >= 6 engineRT = builderRT->buildEngineWithConfig(*networkRT, *configRT); +#else + engineRT = builderRT->buildCudaEngine(*networkRT); +#endif if(engineRT == nullptr) FatalError("cloud not build cuda engine") // we don't need the network any more