dla
This commit is contained in:
+1
-1
@@ -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
@@ -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() {
|
||||||
|
|||||||
@@ -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,
|
||||||
|
|||||||
Reference in New Issue
Block a user