diff --git a/include/tkDNN/Network.h b/include/tkDNN/Network.h index d21659a..f3b623d 100644 --- a/include/tkDNN/Network.h +++ b/include/tkDNN/Network.h @@ -1,6 +1,7 @@ #ifndef NETWORK_H #define NETWORK_H +#include #include "utils.h" namespace tk { namespace dnn { @@ -62,6 +63,8 @@ public: bool fp16, dla, int8; bool dontLoadWeights; + std::string fileImgList; + std::string fileLabelList; }; }} diff --git a/src/Network.cpp b/src/Network.cpp index 24152ec..2e4833b 100644 --- a/src/Network.cpp +++ b/src/Network.cpp @@ -34,6 +34,12 @@ Network::Network(dataDim_t input_dim) { int8 = true; } } + if(const char* env_p = std::getenv("TKDNN_CALIB_IMG_PATH")) + fileImgList = env_p; + + if(const char* env_p = std::getenv("TKDNN_CALIB_LABEL_PATH")) + fileLabelList = env_p; + if(fp16) std::cout<setInt8Mode(true); configRT->setFlag(BuilderFlag::kINT8); BatchStream calibrationStream(dim, 1, 100, //TODO: check if 100 images are sufficient to the calibration (or 4951) - "/home/xavier/Documents/tkDNN/demo/COCO_val2017/all_images.txt", - "/home/xavier/Documents/tkDNN/demo/COCO_val2017/all_labels.txt"); + net->fileImgList, net->fileLabelList); std::string modelName = name; calibrator.reset(new Int8EntropyCalibrator(calibrationStream, 1, "./" + modelName.substr(0, modelName.find('.')) + "-calibration.table",