This commit is contained in:
Roberto Cavicchioli
2019-03-06 16:39:36 +01:00
parent c22219ad16
commit 3714155809
3 changed files with 17 additions and 2 deletions
+1 -1
View File
@@ -59,7 +59,7 @@ public:
dataDim_t input_dim; dataDim_t input_dim;
dataDim_t getOutputDim(); dataDim_t getOutputDim();
bool fp16; bool fp16, dla;
}; };
}} }}
+9 -1
View File
@@ -24,12 +24,20 @@ Network::Network(dataDim_t input_dim) {
num_layers = 0; num_layers = 0;
fp16 = false; fp16 = false;
if(const char* env_p = std::getenv("TKDNN_MODE")) dla = false;
if(const char* env_p = std::getenv("TKDNN_MODE")) {
if(strcmp(env_p, "FP16") == 0) if(strcmp(env_p, "FP16") == 0)
fp16 = true; fp16 = true;
else if(strcmp(env_p, "DLA") == 0) {
dla = true;
fp16 = true;
}
}
if(fp16) if(fp16)
std::cout<<COL_REDB<<"!! FP16 INERENCE ENABLED !!"<<COL_END<<"\n"; std::cout<<COL_REDB<<"!! FP16 INERENCE ENABLED !!"<<COL_END<<"\n";
if(dla)
std::cout<<COL_GREENB<<"!! DLA INERENCE ENABLED !!"<<COL_END<<"\n";
} }
Network::~Network() { Network::~Network() {
+7
View File
@@ -51,6 +51,13 @@ NetworkRT::NetworkRT(Network *net, const char *name) {
dtRT = DataType::kHALF; dtRT = DataType::kHALF;
builderRT->setHalf2Mode(true); builderRT->setHalf2Mode(true);
} }
if(net->dla && builderRT->getNbDLACores() > 0) {
dtRT = DataType::kHALF;
builderRT->setFp16Mode(true);
builderRT->allowGPUFallback(true);
builderRT->setDefaultDeviceType(DeviceType::kDLA);
builderRT->setDLACore(0);
}
//add input layer //add input layer
ITensor *input = networkRT->addInput("data", DataType::kFLOAT, ITensor *input = networkRT->addInput("data", DataType::kFLOAT,