[Benchmark] Add GetModelResoucesNameFromDir APIs (#1640)

This commit is contained in:
DefTruth
2023-03-17 17:05:24 +08:00
committed by GitHub
parent 4d57b62f83
commit 95820a9892
14 changed files with 207 additions and 60 deletions

View File

@@ -17,6 +17,9 @@
#include <unordered_map>
#include "gflags/gflags.h"
#include "fastdeploy/benchmark/utils.h"
#include <sys/types.h>
#include <dirent.h>
#include <cstring>
#ifdef WIN32
static const char sep = '\\';
@@ -91,3 +94,68 @@ static void PrintBenchmarkInfo(std::unordered_map<std::string,
#endif
return;
}
static bool GetModelResoucesNameFromDir(
const std::string& path, std::string* resource_name,
const std::string& suffix = "pdmodel") {
DIR *p_dir;
struct dirent *ptr;
if (!(p_dir = opendir(path.c_str()))) {
return false;
}
bool find = false;
while ((ptr = readdir(p_dir)) != 0) {
if (strcmp(ptr->d_name, ".") != 0 && strcmp(ptr->d_name, "..") != 0) {
std::string tmp_file_name = ptr->d_name;
if (tmp_file_name.find(suffix) != std::string::npos) {
if (suffix == "pdiparams") {
if (tmp_file_name.find("info") == std::string::npos) {
find = true;
*resource_name = tmp_file_name;
break;
}
} else {
find = true;
*resource_name = tmp_file_name;
break;
}
} else {
if (suffix == "yml") {
if (tmp_file_name.find("yaml") != std::string::npos) {
find = true;
*resource_name = tmp_file_name;
break;
}
} else if (suffix == "yaml") {
if (tmp_file_name.find("yml") != std::string::npos) {
find = true;
*resource_name = tmp_file_name;
break;
}
}
}
}
}
closedir(p_dir);
return find;
}
static bool UpdateModelResourceName(
std::string* model_name, std::string* params_name,
std::string* config_name, fastdeploy::ModelFormat* model_format,
std::unordered_map<std::string, std::string>& config_info,
bool use_config_file = true, bool use_quant_model = false) {
*model_format = fastdeploy::ModelFormat::PADDLE;
if (!(GetModelResoucesNameFromDir(FLAGS_model, model_name, "pdmodel")
&& GetModelResoucesNameFromDir(FLAGS_model, params_name, "pdiparams"))) {
std::cout << "Can not find Paddle model resources." << std::endl;
return false;
}
if (use_config_file) {
if (!GetModelResoucesNameFromDir(FLAGS_model, config_name, "yml")) {
std::cout << "Can not find config yaml resources." << std::endl;
return false;
}
}
return true;
}