[TF FE] Enabled MMAP for SavedModel and MetaGraph (#18471)

* Moved mmap to openvino/util as a shared functionality

* Enabled MMAP for SavedModel and MetaGraph

* Fixed CMake

* Fixed a lost line

* Simplified code for compilers

* Aligned with an actual master

* Enabled mmap by default and added test

* Suppressed warning, added test for MetaGraph, additional messages

---------

Co-authored-by: Andrei Kochin <andrei.kochin@intel.com>
This commit is contained in:
Georgy Krivoruchko 2023-08-15 07:36:08 -07:00 committed by GitHub
parent d13ae31a61
commit 18ab677952
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
10 changed files with 123 additions and 33 deletions

View File

@ -189,6 +189,8 @@ bool FrontEnd::supported_impl(const std::vector<ov::Any>& variants) const {
ov::frontend::InputModel::Ptr FrontEnd::load_impl(const std::vector<ov::Any>& variants) const {
// Last boolean flag in `variants` (if presented) is reserved for FE configuration
size_t extra_variants_num = variants.size() > 0 && variants[variants.size() - 1].is<bool>() ? 1 : 0;
// Enable mmap by default
bool mmap_enabled = variants[variants.size() - 1].is<bool>() ? variants[variants.size() - 1].as<bool>() : true;
// For TF1 models it can be a case of two input variants: input model and v1 checkpoints
FRONT_END_GENERAL_CHECK(
@ -203,7 +205,7 @@ ov::frontend::InputModel::Ptr FrontEnd::load_impl(const std::vector<ov::Any>& va
return std::make_shared<InputModel>(std::make_shared<GraphIteratorProto>(model_path), m_telemetry);
} else if (GraphIteratorSavedModel::is_supported(model_path)) {
std::shared_ptr<GraphIteratorSavedModel> graph_iterator;
graph_iterator = std::make_shared<GraphIteratorSavedModel>(model_path, std::string("serve"));
graph_iterator = std::make_shared<GraphIteratorSavedModel>(model_path, std::string("serve"), mmap_enabled);
return std::make_shared<InputModel>(graph_iterator,
m_telemetry,
graph_iterator->get_variables_index(),
@ -212,7 +214,7 @@ ov::frontend::InputModel::Ptr FrontEnd::load_impl(const std::vector<ov::Any>& va
nullptr,
true);
} else if (GraphIteratorMeta::is_supported(model_path)) {
auto graph_iterator = std::make_shared<GraphIteratorMeta>(model_path);
auto graph_iterator = std::make_shared<GraphIteratorMeta>(model_path, mmap_enabled);
return std::make_shared<InputModel>(graph_iterator,
m_telemetry,
graph_iterator->get_variables_index(),
@ -257,7 +259,7 @@ ov::frontend::InputModel::Ptr FrontEnd::load_impl(const std::vector<ov::Any>& va
auto saved_model_tags = paths[1];
if (GraphIteratorSavedModel::is_supported(model_path)) {
std::shared_ptr<GraphIteratorSavedModel> graph_iterator;
graph_iterator = std::make_shared<GraphIteratorSavedModel>(model_path, saved_model_tags);
graph_iterator = std::make_shared<GraphIteratorSavedModel>(model_path, saved_model_tags, mmap_enabled);
return std::make_shared<InputModel>(graph_iterator,
m_telemetry,
graph_iterator->get_variables_index(),
@ -275,7 +277,9 @@ ov::frontend::InputModel::Ptr FrontEnd::load_impl(const std::vector<ov::Any>& va
return std::make_shared<InputModel>(std::make_shared<GraphIteratorProto>(model_path), m_telemetry);
} else if (GraphIteratorSavedModel::is_supported(model_path)) {
std::shared_ptr<GraphIteratorSavedModel> graph_iterator;
graph_iterator = std::make_shared<GraphIteratorSavedModel>(model_path, std::string(META_GRAPH_DEFAULT_TAG));
graph_iterator = std::make_shared<GraphIteratorSavedModel>(model_path,
std::string(META_GRAPH_DEFAULT_TAG),
mmap_enabled);
return std::make_shared<InputModel>(graph_iterator,
m_telemetry,
graph_iterator->get_variables_index(),
@ -284,7 +288,7 @@ ov::frontend::InputModel::Ptr FrontEnd::load_impl(const std::vector<ov::Any>& va
nullptr,
true);
} else if (GraphIteratorMeta::is_supported(model_path)) {
auto graph_iterator = std::make_shared<GraphIteratorMeta>(model_path);
auto graph_iterator = std::make_shared<GraphIteratorMeta>(model_path, mmap_enabled);
return std::make_shared<InputModel>(graph_iterator,
m_telemetry,
graph_iterator->get_variables_index(),
@ -329,7 +333,7 @@ ov::frontend::InputModel::Ptr FrontEnd::load_impl(const std::vector<ov::Any>& va
auto saved_model_tags = ov::util::wstring_to_string(paths[1]);
if (GraphIteratorSavedModel::is_supported(model_path)) {
std::shared_ptr<GraphIteratorSavedModel> graph_iterator;
graph_iterator = std::make_shared<GraphIteratorSavedModel>(model_path, saved_model_tags);
graph_iterator = std::make_shared<GraphIteratorSavedModel>(model_path, saved_model_tags, mmap_enabled);
return std::make_shared<InputModel>(graph_iterator,
m_telemetry,
graph_iterator->get_variables_index(),

View File

@ -31,11 +31,13 @@ class GraphIteratorMeta : public GraphIteratorProto {
std::shared_ptr<VariablesIndex> m_variables_index;
std::shared_ptr<std::map<std::string, std::string>> m_inputs_map;
std::shared_ptr<std::map<std::string, std::string>> m_outputs_map;
bool m_mmap_enabled;
public:
template <typename T>
GraphIteratorMeta(const std::basic_string<T>& path)
: m_metagraph_def(std::make_shared<::tensorflow::MetaGraphDef>()) {
GraphIteratorMeta(const std::basic_string<T>& path, const bool mmap_enabled)
: m_metagraph_def(std::make_shared<::tensorflow::MetaGraphDef>()),
m_mmap_enabled(mmap_enabled) {
this->read_meta(path);
}
@ -75,7 +77,7 @@ private:
std::basic_string<T> varIndexPath = get_variables_index_name<T>(model_path);
if (ov::util::file_exists(varIndexPath)) {
m_variables_index = std::make_shared<VariablesIndex>();
m_variables_index = std::make_shared<VariablesIndex>(m_mmap_enabled);
std::ifstream vi_stream{varIndexPath.c_str(), std::ifstream::in | std::ifstream::binary};
FRONT_END_GENERAL_CHECK(vi_stream && vi_stream.is_open(), "MetaGraph's variable index file does not exist");
FRONT_END_GENERAL_CHECK(m_variables_index->read_variables(vi_stream, model_path, false),

View File

@ -38,11 +38,13 @@ class GraphIteratorSavedModel : public GraphIteratorProto {
std::shared_ptr<VariablesIndex> m_variables_index;
std::shared_ptr<std::map<std::string, std::string>> m_inputs_map;
std::shared_ptr<std::map<std::string, std::string>> m_outputs_map;
bool m_mmap_enabled;
public:
template <typename T>
GraphIteratorSavedModel(const std::basic_string<T>& path, const std::string& tags)
: m_saved_model(std::make_shared<::tensorflow::SavedModel>()) {
GraphIteratorSavedModel(const std::basic_string<T>& path, const std::string& tags, const bool mmap_enabled)
: m_saved_model(std::make_shared<::tensorflow::SavedModel>()),
m_mmap_enabled(mmap_enabled) {
this->read_saved_model(path, tags);
}
@ -74,7 +76,7 @@ private:
std::basic_string<T> varIndexPath = path + get_variables_index_name<T>();
if (ov::util::file_exists(varIndexPath)) {
m_variables_index = std::make_shared<VariablesIndex>();
m_variables_index = std::make_shared<VariablesIndex>(m_mmap_enabled);
std::ifstream vi_stream{varIndexPath.c_str(), std::ifstream::in | std::ifstream::binary};
FRONT_END_GENERAL_CHECK(vi_stream && vi_stream.is_open(),
"[TensorFlow Frontend] Saved Model's variable index file does not exist");

View File

@ -7,7 +7,9 @@
#include "helper_ops/string_constant.hpp"
#include "helper_ops/unsupported_constant.hpp"
#include "input_model.hpp"
#include "ngraph/runtime/shared_buffer.hpp"
#include "openvino/opsets/opset8.hpp"
#include "openvino/util/mmap_object.hpp"
#include "tensor_bundle.pb.h"
using namespace std;
@ -26,22 +28,42 @@ static std::shared_ptr<ov::Node> read_variable(std::shared_ptr<VariablesIndex> v
const ov::Shape shape,
const ::tensorflow::BundleEntryProto& entry,
const NodeContext& node) {
std::vector<T> var_data;
google::protobuf::int64 size = 1;
for (uint64_t i = 0; i < shape.size(); ++i) {
size *= static_cast<google::protobuf::int64>(shape[i]);
}
var_data.resize(size);
TENSORFLOW_OP_VALIDATION(node,
size == static_cast<google::protobuf::int64>(entry.size() / sizeof(T)),
"[TensorFlow Frontend] Internal error: Available data size isn't equal to calculated.");
auto fs = var_index->get_data_file(entry.shard_id());
if (!fs.get()) {
TENSORFLOW_OP_VALIDATION(node, var_index, "[TensorFlow Frontend] Internal error: Cannot get shard file.");
if (var_index->is_mmap_enabled()) {
auto mapped_memory = var_index->get_data_mmap(entry.shard_id());
if (!mapped_memory.get()) {
TENSORFLOW_OP_VALIDATION(node, var_index, "[TensorFlow Frontend] Internal error: Cannot get shard file.");
}
TENSORFLOW_OP_VALIDATION(
node,
static_cast<int64_t>(mapped_memory->size()) >= entry.offset() + entry.size(),
"[TensorFlow Frontend] Internal error: Variable entry size is out of bounds of mapped memory size.");
OPENVINO_SUPPRESS_DEPRECATED_START
return std::make_shared<Constant>(
ov_type,
shape,
std::make_shared<ngraph::runtime::SharedBuffer<std::shared_ptr<MappedMemory>>>(
mapped_memory->data() + entry.offset(),
entry.size(),
mapped_memory));
OPENVINO_SUPPRESS_DEPRECATED_END
} else {
std::vector<T> var_data;
var_data.resize(size);
auto fs = var_index->get_data_file(entry.shard_id());
if (!fs.get()) {
TENSORFLOW_OP_VALIDATION(node, var_index, "[TensorFlow Frontend] Internal error: Cannot get shard file.");
}
fs->seekg(entry.offset(), std::ios::beg);
fs->read(reinterpret_cast<char*>(var_data.data()), entry.size());
return std::make_shared<Constant>(ov_type, shape, var_data);
}
fs->seekg(entry.offset(), std::ios::beg);
fs->read(reinterpret_cast<char*>(var_data.data()), entry.size());
return std::make_shared<Constant>(ov_type, shape, var_data);
}
OutputVector translate_varhandle_op(const NodeContext& node) {

View File

@ -10,6 +10,7 @@
#include "checkpoint_utils.hpp"
#include "graph_iterator_saved_model.hpp"
#include "openvino/core/type/element_type.hpp"
#include "openvino/util/mmap_object.hpp"
#include "tensor_bundle.pb.h"
#include "trackable_object_graph.pb.h"
@ -160,8 +161,13 @@ void VariablesIndex::read_checkpointable_object_graph() {
// It looks like reinterpret_cast artifact
// https://github.com/tensorflow/tensorflow/blob/d90f1947ebcf510b23c238f43c2191e5b3817cb3/tensorflow/cc/experimental/libexport/load.cc#L70
int chg = 6;
shard->second->seekg(entry.offset() + chg);
shard->second->read(data.data(), entry.size() - chg);
if (m_mmap_enabled) {
auto srcPtr = static_cast<char*>(shard->second.mmap->data() + entry.offset() + chg);
std::copy(srcPtr, srcPtr + entry.size() - chg, data.data());
} else {
shard->second.stream->seekg(entry.offset() + chg);
shard->second.stream->read(data.data(), entry.size() - chg);
}
// Might be need to remove this verification:
// https://github.com/tensorflow/tensorflow/blob/d90f1947ebcf510b23c238f43c2191e5b3817cb3/tensorflow/cc/experimental/libexport/load.cc#L73
@ -191,9 +197,14 @@ bool VariablesIndex::read_variables(std::ifstream& vi_stream, const std::string&
} else {
fullPath = path + "." + suffix.data();
}
m_data_files[shard] = std::shared_ptr<std::ifstream>(
new std::ifstream(fullPath.c_str(), std::ifstream::in | std::ifstream::binary));
FRONT_END_GENERAL_CHECK(m_data_files[shard]->is_open(), "Variable index data file does not exist");
if (m_mmap_enabled) {
m_data_files[shard].mmap = load_mmap_object(fullPath);
FRONT_END_GENERAL_CHECK(m_data_files[shard].mmap->data(), "Variable index data cannot be mapped");
} else {
m_data_files[shard].stream = std::shared_ptr<std::ifstream>(
new std::ifstream(fullPath.c_str(), std::ifstream::in | std::ifstream::binary));
FRONT_END_GENERAL_CHECK(m_data_files[shard].stream->is_open(), "Variable index data file does not exist");
}
}
read_checkpointable_object_graph();
@ -215,9 +226,14 @@ bool VariablesIndex::read_variables(std::ifstream& vi_stream, const std::wstring
} else {
fullPath = path + L"." + suffix.data();
}
m_data_files[shard] = std::shared_ptr<std::ifstream>(
new std::ifstream(fullPath.c_str(), std::ifstream::in | std::ifstream::binary));
FRONT_END_GENERAL_CHECK(m_data_files[shard]->is_open(), "Variable index data file does not exist");
if (m_mmap_enabled) {
m_data_files[shard].mmap = load_mmap_object(fullPath);
FRONT_END_GENERAL_CHECK(m_data_files[shard].mmap->data(), L"Variable index data cannot be mapped");
} else {
m_data_files[shard].stream = std::shared_ptr<std::ifstream>(
new std::ifstream(fullPath.c_str(), std::ifstream::in | std::ifstream::binary));
FRONT_END_GENERAL_CHECK(m_data_files[shard].stream->is_open(), L"Variable index data file does not exist");
}
}
read_checkpointable_object_graph();

View File

@ -8,6 +8,7 @@
#include "graph_iterator_proto.hpp"
#include "openvino/util/file_util.hpp"
#include "openvino/util/mmap_object.hpp"
#include "saved_model.pb.h"
namespace ov {
@ -16,6 +17,11 @@ namespace tensorflow {
struct VIBlock;
struct VariableStorage {
std::shared_ptr<std::ifstream> stream;
std::shared_ptr<ov::MappedMemory> mmap;
};
// Stores information about variables index
class VariablesIndex {
// Contains file size for internal checks
@ -25,11 +31,19 @@ class VariablesIndex {
// Contains BundleEntryProto variables list, readed from .index file
std::map<std::string, std::vector<char>> m_variables_index;
// List of opened data files for using with BundleEntryProto
std::map<int32_t, std::shared_ptr<std::ifstream>> m_data_files;
std::map<int32_t, VariableStorage> m_data_files;
// List of mapped variables which could be read using TrackableObjectGraph
std::map<std::string, std::string> m_variables_map;
// Flag shows which file storage is using
bool m_mmap_enabled;
public:
VariablesIndex(bool mmap_enabled = false) : m_mmap_enabled(mmap_enabled) {}
/// \brief Returns mmap_enabled state.
/// \returns True if mmap is enabled, false otherwise
bool is_mmap_enabled(void) const {
return m_mmap_enabled;
}
/// \brief Reads variables from opened variable index file. Can cause an asserts in case of issues.
/// \param vi_stream Opened stream file, file pointer doesn't matter, it will be rewind internally.
/// \param path A path to file with variables data
@ -89,8 +103,20 @@ public:
/// \param shard_id Requested shard_id
/// \returns Valid shared_ptr with ifstream or with nullptr if shard isn't found
std::shared_ptr<std::ifstream> get_data_file(const int32_t shard_id) const {
FRONT_END_GENERAL_CHECK(m_mmap_enabled == false,
"[TensorFlow Frontend] Requested ifstream, but mmap is enabled");
auto result = m_data_files.find(shard_id);
return result != m_data_files.end() ? result->second : nullptr;
return result != m_data_files.end() ? result->second.stream : nullptr;
}
/// \brief Returns shared pointer to a requested shard_id, or nullptr in case of shard_id isn't found
/// \param shard_id Requested shard_id
/// \returns Valid shared_ptr with MappedMemory or with nullptr if shard isn't found
std::shared_ptr<ov::MappedMemory> get_data_mmap(const int32_t shard_id) const {
FRONT_END_GENERAL_CHECK(m_mmap_enabled == true,
"[TensorFlow Frontend] Requested MappedMemory, but mmap is disabled");
auto result = m_data_files.find(shard_id);
return result != m_data_files.end() ? result->second.mmap : nullptr;
}
/// \brief Adds variable mapping to the variables map

View File

@ -158,6 +158,11 @@ TEST_F(FrontEndConversionWithReferenceTestsF, SavedModelWithIntermediateOutput)
}
}
TEST_F(FrontEndConversionWithReferenceTestsF, SavedModelMMAPCompare) {
{ model = convert_model("saved_model_variables"); }
{ model_ref = convert_model("saved_model_variables", nullptr, {}, {}, {}, {}, {}, true); }
}
TEST_F(FrontEndConversionWithReferenceTestsF, SavedModelWithNumericalNames) {
comparator.enable(FunctionsComparator::CmpValues::TENSOR_NAMES);
// The test aims to check that model with only numerical names for operation

View File

@ -550,6 +550,11 @@ TEST_F(FrontEndConversionWithReferenceTestsF, MetaGraphCutIdentity) {
}
}
TEST_F(FrontEndConversionWithReferenceTestsF, MetaGraphMMAPCompare) {
{ model = convert_model("metagraph_variables/graph.meta"); }
{ model_ref = convert_model("metagraph_variables/graph.meta", nullptr, {}, {}, {}, {}, {}, true); }
}
TEST_F(FrontEndConversionWithReferenceTestsF, SplitInFunction) {
{
// create FAKE conversion extension for Split using named ports, this is not required for Split, but it tests

View File

@ -22,7 +22,8 @@ shared_ptr<Model> convert_model(const string& model_path,
const vector<element::Type>& input_types,
const vector<PartialShape>& input_shapes,
const std::vector<std::string>& input_names_to_freeze,
const std::vector<void*>& freeze_values) {
const std::vector<void*>& freeze_values,
const bool disable_mmap) {
FrontEndManager fem;
auto front_end = fem.load_by_framework(TF_FE);
if (!front_end) {
@ -32,7 +33,13 @@ shared_ptr<Model> convert_model(const string& model_path,
front_end->add_extension(conv_ext);
}
auto model_filename = FrontEndTestUtils::make_model_path(string(TEST_TENSORFLOW_MODELS_DIRNAME) + model_path);
auto input_model = front_end->load(model_filename);
ov::frontend::InputModel::Ptr input_model;
if (!disable_mmap) {
input_model = front_end->load(model_filename);
} else {
input_model = front_end->load({model_filename, false});
}
if (!input_model) {
throw "Input model is not read";
}

View File

@ -24,7 +24,8 @@ std::shared_ptr<Model> convert_model(const std::string& model_path,
const std::vector<ov::element::Type>& input_types = {},
const std::vector<ov::PartialShape>& input_shapes = {},
const std::vector<std::string>& input_names_to_freeze = {},
const std::vector<void*>& freeze_values = {});
const std::vector<void*>& freeze_values = {},
const bool disable_mmap = false);
} // namespace tests
} // namespace tensorflow