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

Merge branch 'master' of https://github.com/ceccocats/tkDNN into cnet

Browse Source
This commit is contained in:
Davide Sapienza
2020-01-20 16:41:19 +01:00
parent b23cf9fe70 8a4d1cac17
commit c68d6f318e
6 changed files with 32 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/NetworkRT.cpp
+18 -4
View File
@@ -35,7 +35,7 @@ NetworkRT::NetworkRT(Network *net, const char *name) {
builderRT = createInferBuilder(loggerRT);
std::cout<<"Float16 support: "<<builderRT->platformHasFastFp16()<<"\n";
std::cout<<"Int8 support: "<<builderRT->platformHasFastInt8()<<"\n";
//std::cout<<"DLAs: "<<builderRT->getNbDLACores()<<"\n";
std::cout<<"DLAs: "<<builderRT->getNbDLACores()<<"\n";
networkRT = builderRT->createNetwork();
if(!fileExist(name)) {
@@ -51,7 +51,6 @@ NetworkRT::NetworkRT(Network *net, const char *name) {
dtRT = DataType::kHALF;
builderRT->setHalf2Mode(true);
}
/*
if(net->dla && builderRT->getNbDLACores() > 0) {
dtRT = DataType::kHALF;
builderRT->setFp16Mode(true);
@@ -59,7 +58,6 @@ NetworkRT::NetworkRT(Network *net, const char *name) {
builderRT->setDefaultDeviceType(DeviceType::kDLA);
builderRT->setDLACore(0);
}
*/
//add input layer
ITensor *input = networkRT->addInput("data", DataType::kFLOAT,
@@ -316,10 +314,19 @@ ILayer* NetworkRT::convert_layer(ITensor *input, Activation *l) {
if(l->act_mode == ACTIVATION_LEAKY) {
//std::cout<<"New plugin LEAKY\n";
#if NV_TENSORRT_MAJOR < 6
// plugin version
IPlugin *plugin = new ActivationLeakyRT();
IPluginLayer *lRT = networkRT->addPlugin(&input, 1, *plugin);
checkNULL(lRT);
return lRT;
#else
IActivationLayer *lRT = networkRT->addActivation(*input, ActivationType::kLEAKY_RELU);
lRT->setAlpha(0.1);
checkNULL(lRT);
return lRT;
#endif
} else if(l->act_mode == CUDNN_ACTIVATION_RELU) {
IActivationLayer *lRT = networkRT->addActivation(*input, ActivationType::kRELU);
@@ -391,14 +398,21 @@ ILayer* NetworkRT::convert_layer(ITensor *input, Shortcut *l) {
//std::cout<<"convert Shortcut\n";
//std::cout<<"New plugin Shortcut\n";
ITensor *back_tens = tensors[l->backLayer];
/*
// plugin version
IPlugin *plugin = new ShortcutRT();
ITensor **inputs = new ITensor*[2];
inputs[0] = input;
inputs[1] = back_tens;
IPluginLayer *lRT = networkRT->addPlugin(inputs, 2, *plugin);
checkNULL(lRT);
*/
IElementWiseLayer *lRT = networkRT->addElementWise(*input, *back_tens, ElementWiseOperation::kSUM);
checkNULL(lRT);
return lRT;
}