dla
This commit is contained in:
+9
-1
@@ -24,12 +24,20 @@ Network::Network(dataDim_t input_dim) {
|
||||
num_layers = 0;
|
||||
|
||||
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)
|
||||
fp16 = true;
|
||||
else if(strcmp(env_p, "DLA") == 0) {
|
||||
dla = true;
|
||||
fp16 = true;
|
||||
}
|
||||
}
|
||||
|
||||
if(fp16)
|
||||
std::cout<<COL_REDB<<"!! FP16 INERENCE ENABLED !!"<<COL_END<<"\n";
|
||||
if(dla)
|
||||
std::cout<<COL_GREENB<<"!! DLA INERENCE ENABLED !!"<<COL_END<<"\n";
|
||||
}
|
||||
|
||||
Network::~Network() {
|
||||
|
||||
@@ -51,6 +51,13 @@ 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);
|
||||
builderRT->allowGPUFallback(true);
|
||||
builderRT->setDefaultDeviceType(DeviceType::kDLA);
|
||||
builderRT->setDLACore(0);
|
||||
}
|
||||
|
||||
//add input layer
|
||||
ITensor *input = networkRT->addInput("data", DataType::kFLOAT,
|
||||
|
||||
Reference in New Issue
Block a user