diff --git a/include/onnxruntime/core/graph/graph.h b/include/onnxruntime/core/graph/graph.h index 5359edbbb6..8f81ce14dd 100644 --- a/include/onnxruntime/core/graph/graph.h +++ b/include/onnxruntime/core/graph/graph.h @@ -485,9 +485,12 @@ class Node { #endif static Status LoadFromOrtFormat(const onnxruntime::fbs::Node& fbs_node, Graph& graph, + bool can_use_flatbuffer_for_initializers, const logging::Logger& logger, std::unique_ptr& node); - Status LoadFromOrtFormat(const onnxruntime::fbs::Node& fbs_node, const logging::Logger& logger); + Status LoadFromOrtFormat(const onnxruntime::fbs::Node& fbs_node, + bool can_use_flatbuffer_for_initializers, + const logging::Logger& logger); Status LoadEdgesFromOrtFormat(const onnxruntime::fbs::NodeEdge& fbs_node_edgs, const Graph& graph); /** @@ -1314,17 +1317,18 @@ class Graph { virtual ~Graph(); - static common::Status LoadFromOrtFormat( - const onnxruntime::fbs::Graph& fbs_graph, const Model& owning_model, - const std::unordered_map& domain_to_version, + static common::Status LoadFromOrtFormat(const onnxruntime::fbs::Graph& fbs_graph, const Model& owning_model, + const std::unordered_map& domain_to_version, #if !defined(ORT_MINIMAL_BUILD) - IOnnxRuntimeOpSchemaCollectionPtr schema_registry, + IOnnxRuntimeOpSchemaCollectionPtr schema_registry, #endif - const logging::Logger& logger, std::unique_ptr& graph); + bool can_use_flatbuffer_for_initializers, + const logging::Logger& logger, std::unique_ptr& graph); // deserialize a subgraph static Status LoadFromOrtFormat(const onnxruntime::fbs::Graph& fbs_graph, Graph& parent_graph, const Node& parent_node, + bool can_use_flatbuffer_for_initializers, const logging::Logger& logger, std::unique_ptr& graph); #if !defined(ORT_MINIMAL_BUILD) || defined(ORT_EXTENDED_MINIMAL_BUILD) @@ -1368,7 +1372,8 @@ class Graph { bool strict_shape_type_inference); // Populate Graph instance from ORT format serialized data. - common::Status LoadFromOrtFormat(const onnxruntime::fbs::Graph& fbs_graph); + common::Status LoadFromOrtFormat(const onnxruntime::fbs::Graph& fbs_graph, + bool can_use_flatbuffer_for_initializers); #if !defined(ORT_MINIMAL_BUILD) // Constructor: Given a loaded from model file, construct diff --git a/include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h b/include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h index 3104c4da58..d97f3608cd 100644 --- a/include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h +++ b/include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h @@ -79,6 +79,16 @@ static const char* const kOrtSessionOptionsConfigAllowIntraOpSpinning = "session // has to guarantee that the model bytes are valid until the ORT session using the model bytes is destroyed. static const char* const kOrtSessionOptionsConfigUseORTModelBytesDirectly = "session.use_ort_model_bytes_directly"; +/// +/// Key for using the ORT format model flatbuffer bytes directly for initializers. +/// This avoids copying the bytes and reduces peak memory usage during model loading and initialization. +/// Requires `session.use_ort_model_bytes_directly` to be true. +/// If set, the flatbuffer bytes provided when creating the InferenceSession MUST remain valid for the entire +/// duration of the InferenceSession. +/// +static const char* const kOrtSessionOptionsConfigUseORTModelBytesForInitializers = + "session.use_ort_model_bytes_for_initializers"; + // This should only be specified when exporting an ORT format model for use on a different platform. // If the ORT format model will be used on ARM platforms set to "1". For other platforms set to "0" // Available since version 1.11. diff --git a/onnxruntime/core/framework/session_state_utils.cc b/onnxruntime/core/framework/session_state_utils.cc index 3bb41ad251..a322ae2354 100644 --- a/onnxruntime/core/framework/session_state_utils.cc +++ b/onnxruntime/core/framework/session_state_utils.cc @@ -61,12 +61,15 @@ struct ExtDataValueDeleter { OrtCallback ext_delete_cb; Tensor* p_tensor; void operator()(void*) noexcept { - ext_delete_cb.f(ext_delete_cb.param); + if (ext_delete_cb.f) { + ext_delete_cb.f(ext_delete_cb.param); + } + delete p_tensor; } }; -// given a tensor proto with externdal data return an OrtValue with a tensor for +// given a tensor proto with external data return an OrtValue with a tensor for // that data; the pointers for the tensor data and the tensor itself are owned // by the OrtValue's deleter static inline common::Status ExtDataTensorProtoToTensor(const Env& env, @@ -74,7 +77,6 @@ static inline common::Status ExtDataTensorProtoToTensor(const Env& env, const ONNX_NAMESPACE::TensorProto& tensor_proto, Tensor& tensor, OrtCallback& ext_data_deleter) { ORT_ENFORCE(utils::HasExternalData(tensor_proto)); - ORT_ENFORCE(!proto_path.empty()); void* ext_data_buf = nullptr; SafeInt ext_data_len = 0; @@ -88,8 +90,7 @@ static inline common::Status ExtDataTensorProtoToTensor(const Env& env, std::vector tensor_shape_vec = utils::GetTensorShapeFromTensorProto(tensor_proto); TensorShape tensor_shape{tensor_shape_vec}; - tensor = Tensor(type, tensor_shape, ext_data_buf, OrtMemoryInfo(CPU, - OrtAllocatorType::OrtDeviceAllocator)); + tensor = Tensor(type, tensor_shape, ext_data_buf, OrtMemoryInfo(CPU, OrtAllocatorType::OrtDeviceAllocator)); return common::Status::OK(); } diff --git a/onnxruntime/core/framework/tensor_external_data_info.h b/onnxruntime/core/framework/tensor_external_data_info.h index 3445d8cd67..afc8fda6c3 100644 --- a/onnxruntime/core/framework/tensor_external_data_info.h +++ b/onnxruntime/core/framework/tensor_external_data_info.h @@ -9,20 +9,13 @@ namespace onnxruntime { class ExternalDataInfo { - private: + public: #ifdef _WIN32 using OFFSET_TYPE = int64_t; #else using OFFSET_TYPE = off_t; #endif - std::basic_string rel_path_; - OFFSET_TYPE offset_ = 0; - // 0 means the whole file - size_t length_ = 0; - std::string checksum_; - - public: const std::basic_string& GetRelPath() const { return rel_path_; } OFFSET_TYPE GetOffset() const { return offset_; } @@ -32,7 +25,16 @@ class ExternalDataInfo { // If the value of 'offset' or 'length' field is larger the max value of ssize_t, this function will treat it as a // wrong value and return FAIL. - static common::Status Create(const ::google::protobuf::RepeatedPtrField<::ONNX_NAMESPACE::StringStringEntryProto>& input, - std::unique_ptr& out); + static common::Status Create( + const ::google::protobuf::RepeatedPtrField<::ONNX_NAMESPACE::StringStringEntryProto>& input, + std::unique_ptr& out); + + private: + std::basic_string rel_path_; + OFFSET_TYPE offset_ = 0; + + // 0 means the whole file + size_t length_ = 0; + std::string checksum_; }; -} // namespace onnxruntime \ No newline at end of file +} // namespace onnxruntime diff --git a/onnxruntime/core/framework/tensorprotoutils.cc b/onnxruntime/core/framework/tensorprotoutils.cc index 80a0d17a60..ce384545b8 100644 --- a/onnxruntime/core/framework/tensorprotoutils.cc +++ b/onnxruntime/core/framework/tensorprotoutils.cc @@ -131,20 +131,26 @@ static Status GetExternalDataInfo(const ONNX_NAMESPACE::TensorProto& tensor_prot std::unique_ptr external_data_info; ORT_RETURN_IF_ERROR(onnxruntime::ExternalDataInfo::Create(tensor_proto.external_data(), external_data_info)); - if (tensor_proto_dir != nullptr) { - external_file_path = onnxruntime::ConcatPathComponent(tensor_proto_dir, external_data_info->GetRelPath()); - } else { - external_file_path = external_data_info->GetRelPath(); - } + const auto& location = external_data_info->GetRelPath(); - file_offset = external_data_info->GetOffset(); + if (location == onnxruntime::utils::kTensorProtoMemoryAddressTag) { + external_file_path = location; + } else { + if (tensor_proto_dir != nullptr) { + external_file_path = onnxruntime::ConcatPathComponent(tensor_proto_dir, + external_data_info->GetRelPath()); + } else { + external_file_path = external_data_info->GetRelPath(); + } + } ORT_RETURN_IF_ERROR(onnxruntime::utils::GetSizeInBytesFromTensorProto<0>(tensor_proto, &tensor_byte_size)); const size_t external_data_length = external_data_info->GetLength(); - ORT_RETURN_IF_NOT(external_data_length == 0 || external_data_length == tensor_byte_size, - "TensorProto: ", tensor_proto.name(), " external data size mismatch. Computed size: ", *&tensor_byte_size, - ", external_data.length: ", external_data_length); + "TensorProto: ", tensor_proto.name(), " external data size mismatch. Computed size: ", + *&tensor_byte_size, ", external_data.length: ", external_data_length); + + file_offset = external_data_info->GetOffset(); return Status::OK(); } @@ -597,18 +603,31 @@ Status GetExtDataFromTensorProto(const Env& env, const ORTCHAR_T* model_path, ORT_RETURN_IF_ERROR(GetExternalDataInfo(tensor_proto, t_prot_dir_s, external_data_file_path, file_offset, raw_data_safe_len)); - size_t file_length; - ORT_RETURN_IF_ERROR(env.GetFileLength(external_data_file_path.c_str(), file_length)); + if (external_data_file_path == onnxruntime::utils::kTensorProtoMemoryAddressTag) { + // the value in location is the memory address of the data + ext_data_buf = reinterpret_cast(file_offset); + ext_data_len = raw_data_safe_len; + ext_data_deleter = OrtCallback{nullptr, nullptr}; + } else { + size_t file_length; + // error reporting is inconsistent across platforms. Make sure the full path we attempted to open is included. + auto status = env.GetFileLength(external_data_file_path.c_str(), file_length); + if (!status.IsOK()) { + return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "GetFileLength for ", ToUTF8String(external_data_file_path), + " failed:", status.ErrorMessage()); + } - SafeInt end_of_read(file_offset); - end_of_read += raw_data_safe_len; - ORT_RETURN_IF(file_offset < 0 || end_of_read > gsl::narrow(file_length), + SafeInt end_of_read(file_offset); + end_of_read += raw_data_safe_len; + ORT_RETURN_IF(file_offset < 0 || end_of_read > gsl::narrow(file_length), "External initializer: ", tensor_proto.name(), " offset: ", file_offset, " size to read: ", static_cast(raw_data_safe_len), " given file_length: ", file_length, " are out of bounds or can not be read in full."); - ORT_RETURN_IF_ERROR(GetFileContent(env, external_data_file_path.c_str(), file_offset, raw_data_safe_len, - ext_data_buf, ext_data_deleter)); - ext_data_len = raw_data_safe_len; + ORT_RETURN_IF_ERROR(GetFileContent(env, external_data_file_path.c_str(), file_offset, raw_data_safe_len, + ext_data_buf, ext_data_deleter)); + ext_data_len = raw_data_safe_len; + } + return Status::OK(); } diff --git a/onnxruntime/core/framework/tensorprotoutils.h b/onnxruntime/core/framework/tensorprotoutils.h index d11c7988be..1edfcedc3d 100644 --- a/onnxruntime/core/framework/tensorprotoutils.h +++ b/onnxruntime/core/framework/tensorprotoutils.h @@ -76,6 +76,13 @@ ONNXTensorElementDataType GetTensorElementType(const ONNX_NAMESPACE::TensorProto template common::Status GetSizeInBytesFromTensorProto(const ONNX_NAMESPACE::TensorProto& tensor_proto, size_t* out); +/** +Special marker used to indicate an existing memory buffer contains the TensorProto external data. +If the 'location' field of the external data info is set to this marker, the 'offset' field should contain the +address of the memory containing the data. +*/ +constexpr const ORTCHAR_T* kTensorProtoMemoryAddressTag = ORT_TSTR("*/_ORT_MEM_ADDR_/*"); + // Given a tensor proto with external data obtain a pointer to the data and its length. // The ext_data_deleter argument is updated with a callback that owns/releases the data. common::Status GetExtDataFromTensorProto(const Env& env, const ORTCHAR_T* model_path, diff --git a/onnxruntime/core/graph/graph.cc b/onnxruntime/core/graph/graph.cc index b847e0264d..03078dca09 100644 --- a/onnxruntime/core/graph/graph.cc +++ b/onnxruntime/core/graph/graph.cc @@ -714,12 +714,15 @@ flatbuffers::Offset Node::SaveEdgesToOrtFormat(flatbuffers::FlatB #endif // !defined(ORT_MINIMAL_BUILD) Status Node::LoadFromOrtFormat(const onnxruntime::fbs::Node& fbs_node, Graph& graph, + bool can_use_flatbuffer_for_initializers, const logging::Logger& logger, std::unique_ptr& node) { node = std::make_unique(fbs_node.index(), graph); - return node->LoadFromOrtFormat(fbs_node, logger); + return node->LoadFromOrtFormat(fbs_node, can_use_flatbuffer_for_initializers, logger); } -Status Node::LoadFromOrtFormat(const onnxruntime::fbs::Node& fbs_node, const logging::Logger& logger) { +Status Node::LoadFromOrtFormat(const onnxruntime::fbs::Node& fbs_node, + bool can_use_flatbuffer_for_initializers, + const logging::Logger& logger) { auto LoadNodeArgsFromOrtFormat = [&](const flatbuffers::Vector>* fbs_node_arg_names, std::vector& node_args, @@ -758,7 +761,8 @@ Status Node::LoadFromOrtFormat(const onnxruntime::fbs::Node& fbs_node, const log AttributeProto attr_proto; std::unique_ptr subgraph; ORT_RETURN_IF_ERROR( - fbs::utils::LoadAttributeOrtFormat(*fbs_attr, attr_proto, subgraph, *graph_, *this, logger)); + fbs::utils::LoadAttributeOrtFormat(*fbs_attr, attr_proto, subgraph, *graph_, *this, + can_use_flatbuffer_for_initializers, logger)); // If we have a sub graph in this attributes, it will be loaded into subgraph ptr // while the attribute proto contains the sub graph will have the empty g() field @@ -4187,6 +4191,7 @@ Status Graph::LoadFromOrtFormat(const onnxruntime::fbs::Graph& fbs_graph, #if !defined(ORT_MINIMAL_BUILD) IOnnxRuntimeOpSchemaCollectionPtr schema_registry, #endif + bool can_use_flatbuffer_for_initializers, const logging::Logger& logger, std::unique_ptr& graph) { graph = std::make_unique(owning_model, domain_to_version, #if !defined(ORT_MINIMAL_BUILD) @@ -4196,7 +4201,7 @@ Status Graph::LoadFromOrtFormat(const onnxruntime::fbs::Graph& fbs_graph, // Assume anything in ORT format has already been validated. false); - ORT_RETURN_IF_ERROR(graph->LoadFromOrtFormat(fbs_graph)); + ORT_RETURN_IF_ERROR(graph->LoadFromOrtFormat(fbs_graph, can_use_flatbuffer_for_initializers)); #if !defined(ORT_MINIMAL_BUILD) // in a full build we need to run Resolve to fully populate ResolveContext and Node::op_, @@ -4212,6 +4217,7 @@ Status Graph::LoadFromOrtFormat(const onnxruntime::fbs::Graph& fbs_graph, Status Graph::LoadFromOrtFormat(const onnxruntime::fbs::Graph& fbs_graph, Graph& parent_graph, const Node& parent_node, + bool can_use_flatbuffer_for_initializers, const logging::Logger& logger, std::unique_ptr& graph) { graph = std::make_unique(parent_graph.owning_model_, parent_graph.domain_to_version_, @@ -4223,7 +4229,7 @@ Status Graph::LoadFromOrtFormat(const onnxruntime::fbs::Graph& fbs_graph, // Assume anything in ORT format has already been validated. false); - return graph->LoadFromOrtFormat(fbs_graph); + return graph->LoadFromOrtFormat(fbs_graph, can_use_flatbuffer_for_initializers); } Graph::Graph(const Model& owning_model, @@ -4252,7 +4258,8 @@ Graph::Graph(const Model& owning_model, is_loaded_from_model_file_(true) { // true as the Graph isn't manually constructed from scratch } -common::Status Graph::LoadFromOrtFormat(const onnxruntime::fbs::Graph& fbs_graph) { +common::Status Graph::LoadFromOrtFormat(const onnxruntime::fbs::Graph& fbs_graph, + bool can_use_flatbuffer_for_initializers) { // We deserialize the graph from ORT format in the following order: // 1. Deserialize the initializers and sparse initializers. Convert sparse to dense. // 2. Deserialize the NodeArgs @@ -4282,7 +4289,8 @@ common::Status Graph::LoadFromOrtFormat(const onnxruntime::fbs::Graph& fbs_graph for (const auto* fbs_tensor : *fbs_initializers) { ORT_RETURN_IF(nullptr == fbs_tensor, "Initializer tensor is missing. Invalid ORT format model."); TensorProto* initializer = deserialized_proto_data_.add_initializer(); - ORT_RETURN_IF_ERROR(fbs::utils::LoadInitializerOrtFormat(*fbs_tensor, *initializer)); + ORT_RETURN_IF_ERROR(fbs::utils::LoadInitializerOrtFormat(*fbs_tensor, *initializer, + can_use_flatbuffer_for_initializers)); auto p = name_to_initial_tensor_.emplace(initializer->name(), initializer); if (!p.second) { LOGS(logger_, WARNING) << "Duplicate initializer (dense or ConstantNode): '" << initializer->name() @@ -4342,7 +4350,8 @@ common::Status Graph::LoadFromOrtFormat(const onnxruntime::fbs::Graph& fbs_graph for (const auto* fbs_node : *fbs_nodes) { ORT_RETURN_IF(nullptr == fbs_node, "Node is missing. Invalid ORT format model."); std::unique_ptr node; - ORT_RETURN_IF_ERROR(Node::LoadFromOrtFormat(*fbs_node, *this, logger_, node)); + ORT_RETURN_IF_ERROR(Node::LoadFromOrtFormat(*fbs_node, *this, can_use_flatbuffer_for_initializers, logger_, + node)); ORT_RETURN_IF(node->Index() >= fbs_graph.max_node_index(), "Node index is out of range"); nodes_[node->Index()] = std::move(node); ++num_of_nodes_; diff --git a/onnxruntime/core/graph/graph_flatbuffers_utils.cc b/onnxruntime/core/graph/graph_flatbuffers_utils.cc index d89bae6af2..54e03a985d 100644 --- a/onnxruntime/core/graph/graph_flatbuffers_utils.cc +++ b/onnxruntime/core/graph/graph_flatbuffers_utils.cc @@ -5,6 +5,7 @@ #include "core/flatbuffers/flatbuffers_utils.h" #include "core/flatbuffers/schema/ort.fbs.h" #include "core/framework/tensorprotoutils.h" +#include "core/framework/tensor_external_data_info.h" #include "graph_flatbuffers_utils.h" #include "flatbuffers/flatbuffers.h" @@ -171,8 +172,8 @@ Status SaveAttributeOrtFormat(flatbuffers::FlatBufferBuilder& builder, #endif -Status LoadInitializerOrtFormat(const fbs::Tensor& fbs_tensor, - TensorProto& initializer) { +Status LoadInitializerOrtFormat(const fbs::Tensor& fbs_tensor, TensorProto& initializer, + bool can_use_flatbuffer_for_initializers) { initializer.Clear(); LOAD_STR_FROM_ORT_FORMAT(initializer, name, fbs_tensor.name()); @@ -196,8 +197,29 @@ Status LoadInitializerOrtFormat(const fbs::Tensor& fbs_tensor, const auto* fbs_raw_data = fbs_tensor.raw_data(); ORT_RETURN_IF(nullptr == fbs_raw_data, "Missing raw data for initializer. Invalid ORT format model."); - // fbs_raw_data is uint8_t vector, so the size is byte size - initializer.set_raw_data(fbs_raw_data->Data(), fbs_raw_data->size()); + if (can_use_flatbuffer_for_initializers && fbs_raw_data->size() > 127) { + initializer.set_data_location(ONNX_NAMESPACE::TensorProto_DataLocation_EXTERNAL); + + static_assert(sizeof(void*) <= sizeof(ExternalDataInfo::OFFSET_TYPE)); + const void* data_offset = fbs_raw_data->Data(); + // we reinterpret_cast this back to void* in tensorprotoutils.cc:GetExtDataFromTensorProto. + // use intptr_t as OFFSET_TYPE is signed. in theory you could get a weird looking value if the address uses the + // high bit, but that should be unlikely in a scenario where we care about memory usage enough to use this path. + auto offset = gsl::narrow(reinterpret_cast(data_offset)); + + ONNX_NAMESPACE::StringStringEntryProto* entry = initializer.mutable_external_data()->Add(); + entry->set_key("location"); + entry->set_value(ToUTF8String(onnxruntime::utils::kTensorProtoMemoryAddressTag)); + entry = initializer.mutable_external_data()->Add(); + entry->set_key("offset"); + entry->set_value(std::to_string(offset)); + entry = initializer.mutable_external_data()->Add(); + entry->set_key("length"); + entry->set_value(std::to_string(fbs_raw_data->size())); + } else { + // fbs_raw_data is uint8_t vector, so the size is byte size + initializer.set_raw_data(fbs_raw_data->Data(), fbs_raw_data->size()); + } } return Status::OK(); @@ -231,6 +253,7 @@ Status LoadAttributeOrtFormat(const fbs::Attribute& fbs_attr, ONNX_NAMESPACE::AttributeProto& attr_proto, std::unique_ptr& sub_graph, onnxruntime::Graph& graph, onnxruntime::Node& node, + bool can_use_flatbuffer_for_initializers, const logging::Logger& logger) { attr_proto.Clear(); LOAD_STR_FROM_ORT_FORMAT(attr_proto, name, fbs_attr.name()); @@ -253,7 +276,8 @@ Status LoadAttributeOrtFormat(const fbs::Attribute& fbs_attr, case AttributeProto_AttributeType_TENSOR: { auto fbs_tensor = fbs_attr.t(); ORT_RETURN_IF(nullptr == fbs_tensor, "Null tensor attribute. Invalid ORT format model."); - ORT_RETURN_IF_ERROR(LoadInitializerOrtFormat(*fbs_tensor, *attr_proto.mutable_t())); + ORT_RETURN_IF_ERROR(LoadInitializerOrtFormat(*fbs_tensor, *attr_proto.mutable_t(), + can_use_flatbuffer_for_initializers)); } break; case AttributeProto_AttributeType_GRAPH: { // If the attribute type is a graph, we will create an empty graph in attr_proto so that the ONNX checker @@ -261,7 +285,9 @@ Status LoadAttributeOrtFormat(const fbs::Attribute& fbs_attr, auto fbs_graph = fbs_attr.g(); ORT_RETURN_IF(nullptr == fbs_graph, "Null graph attribute. Invalid ORT format model."); attr_proto.mutable_g()->set_name("Empty graph proto from deserialization of ORT format model"); - ORT_RETURN_IF_ERROR(onnxruntime::Graph::LoadFromOrtFormat(*fbs_graph, graph, node, logger, sub_graph)); + ORT_RETURN_IF_ERROR(onnxruntime::Graph::LoadFromOrtFormat(*fbs_graph, graph, node, + can_use_flatbuffer_for_initializers, + logger, sub_graph)); } break; case AttributeProto_AttributeType_FLOATS: { auto fbs_floats = fbs_attr.floats(); @@ -294,7 +320,8 @@ Status LoadAttributeOrtFormat(const fbs::Attribute& fbs_attr, tensors->Reserve(fbs_tensors->size()); for (const auto* fbs_tensor : *fbs_tensors) { ORT_RETURN_IF(nullptr == fbs_tensor, "Null tensor in tensors attribute. Invalid ORT format model."); - ORT_RETURN_IF_ERROR(LoadInitializerOrtFormat(*fbs_tensor, *tensors->Add())); + ORT_RETURN_IF_ERROR(LoadInitializerOrtFormat(*fbs_tensor, *tensors->Add(), + can_use_flatbuffer_for_initializers)); } } break; diff --git a/onnxruntime/core/graph/graph_flatbuffers_utils.h b/onnxruntime/core/graph/graph_flatbuffers_utils.h index 77d59b45d1..eeb5353dd1 100644 --- a/onnxruntime/core/graph/graph_flatbuffers_utils.h +++ b/onnxruntime/core/graph/graph_flatbuffers_utils.h @@ -49,8 +49,19 @@ onnxruntime::common::Status SaveAttributeOrtFormat( flatbuffers::Offset& fbs_attr, const Path& model_path, const onnxruntime::Graph* subgraph); -onnxruntime::common::Status LoadInitializerOrtFormat( - const fbs::Tensor& fbs_tensor, ONNX_NAMESPACE::TensorProto& initializer); +/// +/// Load an initializer from an ORT format flatbuffer. +/// +/// Flatbuffer Tensor +/// TensorProto to load data into +/// +/// If true, set the TensorProto to point to the memory in the flatbuffer instead of copying data. +/// This requires the buffer to remain valid for the entire duration of the InferenceSession. +/// +/// Status +onnxruntime::common::Status LoadInitializerOrtFormat(const fbs::Tensor& fbs_tensor, + ONNX_NAMESPACE::TensorProto& initializer, + bool can_use_flatbuffer_for_initializers = false); onnxruntime::common::Status LoadSparseInitializerOrtFormat(const fbs::SparseTensor& fbs_sparse_tensor, ONNX_NAMESPACE::SparseTensorProto& initializer); @@ -62,6 +73,7 @@ onnxruntime::common::Status LoadAttributeOrtFormat(const fbs::Attribute& fbs_att ONNX_NAMESPACE::AttributeProto& attr_proto, std::unique_ptr& sub_graph, onnxruntime::Graph& graph, onnxruntime::Node& node, + bool can_use_flatbuffer_for_initializers, const logging::Logger& logger); } // namespace utils diff --git a/onnxruntime/core/graph/model.cc b/onnxruntime/core/graph/model.cc index 12cdaf04e2..75b3749997 100644 --- a/onnxruntime/core/graph/model.cc +++ b/onnxruntime/core/graph/model.cc @@ -770,6 +770,7 @@ common::Status Model::LoadFromOrtFormat(const fbs::Model& fbs_model, #if !defined(ORT_MINIMAL_BUILD) const IOnnxRuntimeOpSchemaRegistryList* local_registries, #endif + bool can_use_flatbuffer_for_initializers, const logging::Logger& logger, std::unique_ptr& model) { model = std::make_unique(); @@ -827,10 +828,11 @@ common::Status Model::LoadFromOrtFormat(const fbs::Model& fbs_model, ORT_RETURN_IF(nullptr == fbs_graph, "Graph is null. Invalid ORT format model."); #if !defined(ORT_MINIMAL_BUILD) - ORT_RETURN_IF_ERROR(Graph::LoadFromOrtFormat(*fbs_graph, *model, domain_to_version, schema_registry, logger, - model->graph_)); + ORT_RETURN_IF_ERROR(Graph::LoadFromOrtFormat(*fbs_graph, *model, domain_to_version, schema_registry, + can_use_flatbuffer_for_initializers, logger, model->graph_)); #else - ORT_RETURN_IF_ERROR(Graph::LoadFromOrtFormat(*fbs_graph, *model, domain_to_version, logger, model->graph_)); + ORT_RETURN_IF_ERROR(Graph::LoadFromOrtFormat(*fbs_graph, *model, domain_to_version, + can_use_flatbuffer_for_initializers, logger, model->graph_)); #endif return Status::OK(); } diff --git a/onnxruntime/core/graph/model.h b/onnxruntime/core/graph/model.h index 14939eead3..cf9a5c66cc 100644 --- a/onnxruntime/core/graph/model.h +++ b/onnxruntime/core/graph/model.h @@ -40,7 +40,8 @@ struct ModelOptions { bool strict_shape_type_inference; ModelOptions(bool allow_released_opsets_only, bool strict_shape_type_inference) - : allow_released_opsets_only(allow_released_opsets_only), strict_shape_type_inference(strict_shape_type_inference) {} + : allow_released_opsets_only(allow_released_opsets_only), + strict_shape_type_inference(strict_shape_type_inference) {} ModelOptions() : ModelOptions(true, false) {} }; @@ -297,6 +298,7 @@ class Model { #if !defined(ORT_MINIMAL_BUILD) const IOnnxRuntimeOpSchemaRegistryList* local_registries, #endif + bool can_use_flatbuffer_for_initializers, const logging::Logger& logger, std::unique_ptr& model); diff --git a/onnxruntime/core/session/inference_session.cc b/onnxruntime/core/session/inference_session.cc index 2eced63620..1c513ad265 100644 --- a/onnxruntime/core/session/inference_session.cc +++ b/onnxruntime/core/session/inference_session.cc @@ -972,9 +972,11 @@ Status InferenceSession::LoadOrtModel(const PathString& model_uri) { Status InferenceSession::LoadOrtModel(const void* model_data, int model_data_len) { return LoadOrtModelWithLoader([&]() { + const auto& config_options = GetSessionOptions().config_options; const auto use_ort_model_bytes_directly = - GetSessionOptions().config_options.GetConfigOrDefault(kOrtSessionOptionsConfigUseORTModelBytesDirectly, "0"); - if (use_ort_model_bytes_directly != "1") { + config_options.GetConfigOrDefault(kOrtSessionOptionsConfigUseORTModelBytesDirectly, "0") == "1"; + + if (!use_ort_model_bytes_directly) { // copy bytes as we need them to be available when InferenceSession::Initialize is called later. ort_format_model_bytes_data_holder_.resize(model_data_len); std::copy_n(reinterpret_cast(model_data), model_data_len, @@ -1025,15 +1027,25 @@ Status InferenceSession::LoadOrtModelWithLoader(std::function load_ort const auto* fbs_model = fbs_session->model(); ORT_RETURN_IF(nullptr == fbs_model, "Missing Model. Invalid ORT format model."); + // if we're using the bytes directly because kOrtSessionOptionsConfigUseORTModelBytesDirectly was set and the user + // provided an existing buffer of bytes when creating the InferenceSession, ort_format_model_bytes_data_holder_ + // will be empty. + // if that is the case we also allow creating initializers that directly use those bytes. + const auto& config_options = session_options_.config_options; + using_ort_model_bytes_for_initializers_ = + ort_format_model_bytes_data_holder_.empty() && + config_options.GetConfigOrDefault(kOrtSessionOptionsConfigUseORTModelBytesForInitializers, "0") == "1"; + // need to go from unique_ptr to shared_ptr when moving into model_ std::unique_ptr tmp_model; #if !defined(ORT_MINIMAL_BUILD) ORT_RETURN_IF_ERROR(Model::LoadFromOrtFormat(*fbs_model, HasLocalSchema() ? &custom_schema_registries_ : nullptr, + using_ort_model_bytes_for_initializers_, *session_logger_, tmp_model)); - #else - ORT_RETURN_IF_ERROR(Model::LoadFromOrtFormat(*fbs_model, *session_logger_, tmp_model)); + ORT_RETURN_IF_ERROR(Model::LoadFromOrtFormat(*fbs_model, using_ort_model_bytes_for_initializers_, *session_logger_, + tmp_model)); #endif ORT_RETURN_IF_ERROR(SaveModelMetadata(*tmp_model)); @@ -1264,9 +1276,9 @@ common::Status InferenceSession::Initialize() { // // NOTE: UpdateProvidersWithSharedAllocators is replace-only and will not insert a new allocator into the EP, so // it must be called after RegisterAllocator. - std::string use_env_allocators = session_options_.config_options.GetConfigOrDefault(kOrtSessionOptionsConfigUseEnvAllocators, - "0"); - if (use_env_allocators == "1") { + bool use_env_allocators = + session_options_.config_options.GetConfigOrDefault(kOrtSessionOptionsConfigUseEnvAllocators, "0") == "1"; + if (use_env_allocators) { LOGS(*session_logger_, INFO) << "This session will use the allocator registered with the environment."; UpdateProvidersWithSharedAllocators(); } @@ -1446,10 +1458,10 @@ common::Status InferenceSession::Initialize() { is_inited_ = true; - // we don't directly use the ORT format bytes currently, so free those now - // TODO, we may need to keep the bytes if we are using the offset directly in the initializers - ort_format_model_bytes_ = gsl::span(); - std::vector().swap(ort_format_model_bytes_data_holder_); + if (!using_ort_model_bytes_for_initializers_) { + ort_format_model_bytes_ = gsl::span(); + std::vector().swap(ort_format_model_bytes_data_holder_); + } // and log telemetry bool model_has_fp16_inputs = ModelHasFP16Inputs(graph); diff --git a/onnxruntime/core/session/inference_session.h b/onnxruntime/core/session/inference_session.h index 6676b521a3..aae2394ebd 100644 --- a/onnxruntime/core/session/inference_session.h +++ b/onnxruntime/core/session/inference_session.h @@ -791,6 +791,8 @@ class InferenceSession { // "session.use_ort_model_bytes_directly" to "1", this will be empty std::vector ort_format_model_bytes_data_holder_; + bool using_ort_model_bytes_for_initializers_{false}; + // Container to store pre-packed weights to share between sessions. // The life-cycle of the cache itself is maintained by the user and the user will ensure // the cache is valid until any session reliant on it is still in scope. diff --git a/onnxruntime/test/framework/ort_model_only_test.cc b/onnxruntime/test/framework/ort_model_only_test.cc index daadfa9956..a73780c4ee 100644 --- a/onnxruntime/test/framework/ort_model_only_test.cc +++ b/onnxruntime/test/framework/ort_model_only_test.cc @@ -36,6 +36,7 @@ struct OrtModelTestInfo { std::vector> configs; bool run_use_buffer{false}; bool disable_copy_ort_buffer{false}; + bool use_buffer_for_initializers{false}; }; static void RunOrtModel(const OrtModelTestInfo& test_info) { @@ -47,6 +48,10 @@ static void RunOrtModel(const OrtModelTestInfo& test_info) { if (test_info.disable_copy_ort_buffer) { ASSERT_STATUS_OK(so.config_options.AddConfigEntry(kOrtSessionOptionsConfigUseORTModelBytesDirectly, "1")); + + if (test_info.use_buffer_for_initializers) { + ASSERT_STATUS_OK(so.config_options.AddConfigEntry(kOrtSessionOptionsConfigUseORTModelBytesForInitializers, "1")); + } } std::vector model_data; @@ -460,6 +465,15 @@ TEST(OrtModelOnlyTests, LoadOrtFormatModelFromBufferNoCopy) { RunOrtModel(test_info); } +// Load the model from a buffer instead of a file path, and not copy the buffer in session creation +TEST(OrtModelOnlyTests, LoadOrtFormatModelFromBufferNoCopyInitializersUseBuffer) { + OrtModelTestInfo test_info = GetTestInfoForLoadOrtFormatModel(); + test_info.run_use_buffer = true; + test_info.disable_copy_ort_buffer = true; + test_info.use_buffer_for_initializers = true; + RunOrtModel(test_info); +} + #if !defined(DISABLE_ML_OPS) // test that we can deserialize and run a previously saved ORT format model // for a model with sequence and map outputs @@ -543,6 +557,5 @@ TEST(OrtModelOnlyTests, TestBackwardsCompat) { } #endif // !defined(DISABLE_ML_OPS) - } // namespace test } // namespace onnxruntime