fz-mirrors/MaaAssistantArknights
0
0
Fork 0
mirror of https://github.com/MaaAssistantArknights/MaaAssistantArknights.git synced 2026-07-19 10:32:19 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat: 惰性加载onnx模型

Browse Source
This commit is contained in:
MistEO
2023-04-07 22:44:19 +08:00
parent 764a7cd8fa
commit da0b860f4a
2 changed files with 20 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

22
src/MaaCore/Config/OnnxSessions.cpp
View File

@@ -9,12 +9,26 @@
bool asst::OnnxSessions::load(const std::filesystem::path& path)
{
LogTraceFunction;
Log.info("load", path);
Log.info("record path", path);
if (!std::filesystem::exists(path)) {
Log.error("file not exist:", path);
return false;
}
std::string name = utils::path_to_utf8_string(path.stem());
Ort::Session session(m_env, path.c_str(), m_options);
m_sessions.insert_or_assign(std::move(name), std::move(session));
m_model_paths.insert_or_assign(name, path);
return true;
}
Ort::Session& asst::OnnxSessions::get(const std::string& key)
{
if (m_sessions.find(key) == m_sessions.end()) {
const auto& path = m_model_paths.at(key);
LogTraceScope(std::string(__FUNCTION__) + " | lazy load: " + utils::path_to_utf8_string(path) + " as " + key);
Ort::Session session(m_env, m_model_paths.at(key).c_str(), m_options);
m_sessions.emplace(key, std::move(session));
}
return m_sessions.at(key);
}

3
src/MaaCore/Config/OnnxSessions.h
View File

@@ -14,11 +14,12 @@ namespace asst
virtual ~OnnxSessions() override = default;
virtual bool load(const std::filesystem::path& path) override;
Ort::Session& get(const std::string& key) { return m_sessions.at(key); }
Ort::Session& get(const std::string& key);
private:
Ort::Env m_env;
Ort::SessionOptions m_options;
std::unordered_map<std::string, Ort::Session> m_sessions;
std::unordered_map<std::string, std::filesystem::path> m_model_paths;
};
}