diff --git a/include/onnxruntime/core/graph/graph.h b/include/onnxruntime/core/graph/graph.h index 360ea68380..46a485c1fc 100644 --- a/include/onnxruntime/core/graph/graph.h +++ b/include/onnxruntime/core/graph/graph.h @@ -32,6 +32,7 @@ #include "core/common/common.h" #include "core/common/const_pointer_container.h" +#include "core/common/inlined_containers_fwd.h" #include "core/common/path.h" #include "core/common/status.h" #include "core/common/logging/logging.h" @@ -675,7 +676,15 @@ class Graph { Note: This currently has linear time complexity. There is room for improvement but it would likely require changes to how initializer tensors are stored and tracked. */ - common::Status ReplaceInitializedTensor(const ONNX_NAMESPACE::TensorProto& new_initializer); + common::Status ReplaceInitializedTensor(ONNX_NAMESPACE::TensorProto new_initializer); + +#if !defined(DISABLE_EXTERNAL_INITIALIZERS) + /** This function takes externally provided data for initializers with external data + * and replaces graph initializers with its content. + */ + common::Status InjectExternalInitializedTensors(const InlinedHashMap& external_initializers); +#endif // !defined(DISABLE_EXTERNAL_INITIALIZERS) + #endif // !defined(ORT_MINIMAL_BUILD) #if !defined(ORT_MINIMAL_BUILD) || defined(ORT_EXTENDED_MINIMAL_BUILD) @@ -1465,6 +1474,9 @@ class Graph { // so they can be used to resolve outer scope dependencies when running BuildConnections for the subgraphs. common::Status SetOuterScopeNodeArgs(const std::unordered_set& outer_scope_node_args); + // Implementation for initializer replacement + Status ReplaceInitializedTensorImpl(ONNX_NAMESPACE::TensorProto new_initializer, bool is_external); + // Clear all unused initializers and NodeArgs void CleanUnusedInitializersAndNodeArgs(const std::unordered_set* initializer_names_to_preserve = nullptr); diff --git a/include/onnxruntime/core/session/onnxruntime_c_api.h b/include/onnxruntime/core/session/onnxruntime_c_api.h index fc8ca5276c..922adca4ab 100644 --- a/include/onnxruntime/core/session/onnxruntime_c_api.h +++ b/include/onnxruntime/core/session/onnxruntime_c_api.h @@ -30,7 +30,7 @@ * * This value is used by some API functions to behave as this version of the header expects. */ -#define ORT_API_VERSION 12 +#define ORT_API_VERSION 12 #ifdef __cplusplus extern "C" { @@ -483,7 +483,7 @@ typedef struct OrtTensorRTProviderOptions { const char* trt_engine_decryption_lib_path; // specify engine decryption library path int trt_force_sequential_engine_build; // force building TensorRT engine sequentially. Default 0 = false, nonzero = true // This is the legacy struct and don't add new fields here. - // For new field that can be represented by string, please add it in include/onnxruntime/core/providers/tensorrt/tensorrt_provider_options.h + // For new field that can be represented by string, please add it in include/onnxruntime/core/providers/tensorrt/tensorrt_provider_options.h // For non-string field, need to create a new separate api to handle it. } OrtTensorRTProviderOptions; @@ -492,9 +492,9 @@ typedef struct OrtTensorRTProviderOptions { * \see OrtApi::SessionOptionsAppendExecutionProvider_MIGraphX */ typedef struct OrtMIGraphXProviderOptions { - int device_id; // hip device id. - int migraphx_fp16_enable; // enable MIGraphX FP16 precision. Default 0 = false, nonzero = true - int migraphx_int8_enable; // enable MIGraphX INT8 precision. Default 0 = false, nonzero = true + int device_id; // hip device id. + int migraphx_fp16_enable; // enable MIGraphX FP16 precision. Default 0 = false, nonzero = true + int migraphx_int8_enable; // enable MIGraphX INT8 precision. Default 0 = false, nonzero = true } OrtMIGraphXProviderOptions; /** \brief OpenVINO Provider Options @@ -3304,6 +3304,31 @@ struct OrtApi { */ ORT_API2_STATUS(SessionOptionsAppendExecutionProvider_MIGraphX, _In_ OrtSessionOptions* options, _In_ const OrtMIGraphXProviderOptions* migraphx_options); + + /** \brief Replace initialized Tensors with external data with the data provided in initializers. + * + * The function will find the initialized TensorProtos with external data in the graph with the provided names and + * replace them with the provided tensors. The API verifies that the TensorProto being replaced + * has an external data reference and has the same name, dimensions and data type as its replacement. The replacement + * will occur before any of the optimizations take place. The data will be copied into the graph + * since TensorProto can't refer to the user provided buffers. + * + * Once the model has been loaded, the OrtValue(s) added to SessionOptions instance will be removed + * from the internal SessionOptions copy to save memory, the user provided buffers can then be deallocated + * and the SessionOptions instance that refers to them can be destroyed. + * + * \param[in] options + * \param[in] initializer_names Array of null terminated UTF-8 encoded strings of the initializers names. + * \param[in] initializers Array of ::OrtValue type + * \param[in] initializers_num Number of elements in the initializer_names and initializers + * + * \snippet{doc} snippets.dox OrtStatus Return Value + * + * \since Version 1.12. + */ + ORT_API2_STATUS(AddExternalInitializers, _In_ OrtSessionOptions* options, + _In_reads_(input_len) const char* const* initializer_names, + _In_reads_(input_len) const OrtValue* const* initializers, size_t initializers_num); }; /* @@ -3373,8 +3398,6 @@ ORT_API_STATUS(OrtSessionOptionsAppendExecutionProvider_CUDA, _In_ OrtSessionOpt * \param device_id HIP device id, starts from zero. */ ORT_API_STATUS(OrtSessionOptionsAppendExecutionProvider_MIGraphX, _In_ OrtSessionOptions* options, int device_id); - - #ifdef __cplusplus } #endif diff --git a/include/onnxruntime/core/session/onnxruntime_cxx_api.h b/include/onnxruntime/core/session/onnxruntime_cxx_api.h index 5c29a0f82a..6f19d2dcb3 100644 --- a/include/onnxruntime/core/session/onnxruntime_cxx_api.h +++ b/include/onnxruntime/core/session/onnxruntime_cxx_api.h @@ -359,6 +359,7 @@ struct SessionOptions : Base { SessionOptions& AddConfigEntry(const char* config_key, const char* config_value); ///< Wraps OrtApi::AddSessionConfigEntry SessionOptions& AddInitializer(const char* name, const OrtValue* ort_val); ///< Wraps OrtApi::AddInitializer + SessionOptions& AddExternalInitializers(const std::vector& names, const std::vector& ort_values); ///< Wraps OrtApi::AddExternalInitializers SessionOptions& AppendExecutionProvider_CUDA(const OrtCUDAProviderOptions& provider_options); ///< Wraps OrtApi::SessionOptionsAppendExecutionProvider_CUDA SessionOptions& AppendExecutionProvider_ROCM(const OrtROCMProviderOptions& provider_options); ///< Wraps OrtApi::SessionOptionsAppendExecutionProvider_ROCM diff --git a/include/onnxruntime/core/session/onnxruntime_cxx_inline.h b/include/onnxruntime/core/session/onnxruntime_cxx_inline.h index 063acb1702..57381301fc 100644 --- a/include/onnxruntime/core/session/onnxruntime_cxx_inline.h +++ b/include/onnxruntime/core/session/onnxruntime_cxx_inline.h @@ -503,6 +503,24 @@ inline SessionOptions& SessionOptions::AddInitializer(const char* name, const Or return *this; } +inline SessionOptions& SessionOptions::AddExternalInitializers(const std::vector& names, + const std::vector& ort_values) { + const size_t inputs_num = names.size(); + if (inputs_num != ort_values.size()) { + ORT_CXX_API_THROW("Expecting names and ort_values to have the same length", ORT_INVALID_ARGUMENT); + } + std::vector names_ptr; + std::vector ort_values_ptrs; + names_ptr.reserve(inputs_num); + ort_values_ptrs.reserve(inputs_num); + for (size_t i = 0; i < inputs_num; ++i) { + names_ptr.push_back(names[i].c_str()); + ort_values_ptrs.push_back(ort_values[i]); + } + ThrowOnError(GetApi().AddExternalInitializers(p_, names_ptr.data(), ort_values_ptrs.data(), inputs_num)); + return *this; +} + inline SessionOptions& SessionOptions::AppendExecutionProvider_CUDA(const OrtCUDAProviderOptions& provider_options) { ThrowOnError(GetApi().SessionOptionsAppendExecutionProvider_CUDA(p_, &provider_options)); return *this; @@ -1133,7 +1151,7 @@ inline T* CustomOpApi::GetTensorMutableData(_Inout_ OrtValue* value) { inline const OrtMemoryInfo* CustomOpApi::GetTensorMemoryInfo(_In_ const OrtValue* value) { const OrtMemoryInfo* mem_info; - ThrowOnError(api_.GetTensorMemoryInfo(value, &mem_info)); + ThrowOnError(api_.GetTensorMemoryInfo(value, &mem_info)); return mem_info; } diff --git a/onnxruntime/core/framework/session_options.cc b/onnxruntime/core/framework/session_options.cc index 2e41b05440..cb625e7f55 100644 --- a/onnxruntime/core/framework/session_options.cc +++ b/onnxruntime/core/framework/session_options.cc @@ -7,30 +7,54 @@ namespace onnxruntime { -Status SessionOptions::AddInitializer(_In_z_ const char* name, _In_ const OrtValue* val) noexcept { - // input validation +namespace { + +Status CheckInitializer(const char* name, const OrtValue* val) { if (name == nullptr) { - return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Received nullptr for name."); + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Received nullptr for name"); } if (val == nullptr) { - return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Received nullptr for OrtValue."); + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Received nullptr for OrtValue"); } if (!val->IsTensor()) { return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Received OrtValue is not a tensor. Only tensors are supported."); } - if (val->Get().OwnsBuffer()) { return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Buffer containing the initializer must be owned by the user."); } + return Status::OK(); +} +} // namespace + +Status SessionOptions::AddInitializer(_In_z_ const char* name, _In_ const OrtValue* val) { + // input validation + ORT_RETURN_IF_ERROR(CheckInitializer(name, val)); // now do the actual work - auto rc = initializers_to_share_map.insert({name, val}); - if (!rc.second) { - return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "An OrtValue for this name has already been added."); + bool result = initializers_to_share_map.emplace(name, val).second; + + if (!result) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "An OrtValue for this name has already been added: ", name); } return Status::OK(); } + +#if !defined(ORT_MINIMAL_BUILD) && !defined(DISABLE_EXTERNAL_INITIALIZERS) +Status SessionOptions::AddExternalInitializers(gsl::span names, gsl::span values) { + const auto init_num = names.size(); + ORT_ENFORCE(init_num == values.size(), "Expecting same size spans"); + external_initializers.reserve(external_initializers.size() + init_num); + for (size_t i = 0; i < init_num; ++i) { + ORT_RETURN_IF_ERROR(CheckInitializer(names[i].c_str(), &values[i])); + bool result = external_initializers.emplace(names[i], values[i]).second; + if (!result) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "An OrtValue for this name has already been added: ", names[i]); + } + } + return Status::OK(); +} +#endif // !defined(ORT_MINIMAL_BUILD) && !defined(DISABLE_EXTERNAL_INITIALIZERS) } // namespace onnxruntime diff --git a/onnxruntime/core/framework/session_options.h b/onnxruntime/core/framework/session_options.h index 09e679a657..c3012299a3 100644 --- a/onnxruntime/core/framework/session_options.h +++ b/onnxruntime/core/framework/session_options.h @@ -6,6 +6,7 @@ #include #include #include "core/common/gsl_suppress.h" +#include "core/common/inlined_containers.h" #include "core/session/onnxruntime_c_api.h" #include "core/optimizer/graph_transformer_level.h" #include "core/util/thread_utils.h" @@ -119,7 +120,13 @@ struct SessionOptions { std::unordered_map initializers_to_share_map; // See onnxruntime_c_api.h for detailed documentation. - Status AddInitializer(_In_z_ const char* name, _In_ const OrtValue* val) noexcept; + Status AddInitializer(_In_z_ const char* name, _In_ const OrtValue* val); + +#if !defined(ORT_MINIMAL_BUILD) && !defined(DISABLE_EXTERNAL_INITIALIZERS) + // Customer supplied pre-processed data for external initializers + InlinedHashMap external_initializers; + Status AddExternalInitializers(gsl::span names, gsl::span values); +#endif // custom function callback to create a thread OrtCustomCreateThreadFn custom_create_thread_fn = nullptr; diff --git a/onnxruntime/core/framework/tensorprotoutils.cc b/onnxruntime/core/framework/tensorprotoutils.cc index 5b5b542643..986cd6b05c 100644 --- a/onnxruntime/core/framework/tensorprotoutils.cc +++ b/onnxruntime/core/framework/tensorprotoutils.cc @@ -125,8 +125,8 @@ static Status GetExternalDataInfo(const ONNX_NAMESPACE::TensorProto& tensor_prot ORT_RETURN_IF_NOT(onnxruntime::utils::HasExternalData(tensor_proto), "Tensor does not have external data to read from."); - ORT_RETURN_IF_NOT(tensor_proto.data_type() != ONNX_NAMESPACE::TensorProto_DataType_STRING, - "External data type cannot be UNDEFINED or STRING."); + ORT_RETURN_IF(!onnxruntime::utils::HasDataType(tensor_proto) || onnxruntime::utils::HasString(tensor_proto), + "External data type cannot be UNDEFINED or STRING."); std::unique_ptr external_data_info; ORT_RETURN_IF_ERROR(onnxruntime::ExternalDataInfo::Create(tensor_proto.external_data(), external_data_info)); @@ -143,7 +143,7 @@ static Status GetExternalDataInfo(const ONNX_NAMESPACE::TensorProto& tensor_prot 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 external data size mismatch. Computed size: ", *&tensor_byte_size, + "TensorProto: ", tensor_proto.name(), " external data size mismatch. Computed size: ", *&tensor_byte_size, ", external_data.length: ", external_data_length); return Status::OK(); @@ -629,6 +629,16 @@ Status TensorProtoToTensor(const Env& env, const ORTCHAR_T* model_path, tensor_proto_dir.size() == 0 ? nullptr : tensor_proto_dir.c_str(), external_data_file_path, file_offset, raw_data_len)); + size_t file_length; + ORT_RETURN_IF_ERROR(env.GetFileLength(external_data_file_path.c_str(), file_length)); + + SafeInt end_of_read(file_offset); + end_of_read += raw_data_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_len), " given file_length: ", file_length, + " are out of bounds or can not be read in full."); + // load the file ORT_RETURN_IF_ERROR(GetFileContent( env, external_data_file_path.c_str(), file_offset, raw_data_len, diff --git a/onnxruntime/core/framework/tensorprotoutils.h b/onnxruntime/core/framework/tensorprotoutils.h index 8b8a3d6ed0..bb60677297 100644 --- a/onnxruntime/core/framework/tensorprotoutils.h +++ b/onnxruntime/core/framework/tensorprotoutils.h @@ -294,6 +294,10 @@ inline bool HasDataType(const ONNX_NAMESPACE::TensorProto& ten_proto) { return ten_proto.data_type() != ONNX_NAMESPACE::TensorProto::UNDEFINED; } +inline bool HasString(const ONNX_NAMESPACE::TensorProto& ten_proto) { + return ten_proto.data_type() == ONNX_NAMESPACE::TensorProto_DataType::TensorProto_DataType_STRING; +} + #ifndef SHARED_PROVIDER inline bool HasName(const ONNX_NAMESPACE::TensorProto& ten_proto) { return ten_proto.has_name(); // XXX diff --git a/onnxruntime/core/graph/graph.cc b/onnxruntime/core/graph/graph.cc index 6e56dffad3..a5a0bc67e8 100644 --- a/onnxruntime/core/graph/graph.cc +++ b/onnxruntime/core/graph/graph.cc @@ -12,6 +12,7 @@ #include "gsl/gsl" #include "core/common/logging/logging.h" +#include "core/common/inlined_containers.h" #include "core/flatbuffers/flatbuffers_utils.h" #include "core/flatbuffers/schema/ort.fbs.h" #include "core/framework/tensor_shape.h" @@ -2885,7 +2886,7 @@ void Graph::RemoveInitializedTensor(const std::string& tensor_name) { } #if !defined(ORT_MINIMAL_BUILD) -Status Graph::ReplaceInitializedTensor(const ONNX_NAMESPACE::TensorProto& new_initializer) { +Status Graph::ReplaceInitializedTensorImpl(ONNX_NAMESPACE::TensorProto new_initializer, bool is_external) { // name_to_initial_tensor_ maps from name to const TensorProto*, so we first // look up the const pointer by name, then find and modify the mutable // pointed-to TensorProto in graph_proto_. @@ -2904,6 +2905,8 @@ Status Graph::ReplaceInitializedTensor(const ONNX_NAMESPACE::TensorProto& new_in return true; }; + ORT_RETURN_IF_NOT(!is_external || utils::HasExternalData(old_initializer), "Trying to replace non-external initializer with external data"); + ORT_RETURN_IF_NOT(dims_eq(), "Replacement tensor's dimensions do not match."); ORT_RETURN_IF_NOT(old_initializer.data_type() == new_initializer.data_type(), "Replacement tensor's data type does not match."); @@ -2917,10 +2920,28 @@ Status Graph::ReplaceInitializedTensor(const ONNX_NAMESPACE::TensorProto& new_in ORT_ENFORCE(existing_entry != mutable_initializers.pointer_end(), "graph_proto_ is not in sync with name_to_initial_tensor_"); - **existing_entry = new_initializer; + **existing_entry = std::move(new_initializer); return Status::OK(); } + +Status Graph::ReplaceInitializedTensor(ONNX_NAMESPACE::TensorProto new_initializer) { + return ReplaceInitializedTensorImpl(std::move(new_initializer), false); +} + +#if !defined(DISABLE_EXTERNAL_INITIALIZERS) +Status Graph::InjectExternalInitializedTensors(const InlinedHashMap& external_initializers) { + for (const auto& e : external_initializers) { + const auto& name = e.first; + const OrtValue& ort_value = e.second; + auto tensor_proto = utils::TensorToTensorProto(ort_value.Get(), name); + ORT_RETURN_IF_ERROR(ReplaceInitializedTensorImpl(std::move(tensor_proto), true)); + LOGS(logger_, INFO) << "Replaced external initializer: " << name; + } + return Status::OK(); +} +#endif // DISABLE_EXTERNAL_INITIALIZERS + #endif // !defined(ORT_MINIMAL_BUILD) bool Graph::GetInitializedTensor(const std::string& tensor_name, const TensorProto*& value) const { diff --git a/onnxruntime/core/optimizer/nchwc_transformer.cc b/onnxruntime/core/optimizer/nchwc_transformer.cc index 57eb1e9ef1..6a41d7b3e3 100644 --- a/onnxruntime/core/optimizer/nchwc_transformer.cc +++ b/onnxruntime/core/optimizer/nchwc_transformer.cc @@ -409,7 +409,7 @@ void NchwcTransformerImpl::TransformConv(Node& node) { nchwc_conv_W_arg = filters_it->second; } else { Initializer conv_W{*conv_W_tensor_proto, graph_.ModelPath()}; - const auto& conv_W_dims = conv_W.dims(); + const auto conv_W_dims = conv_W.dims(); int64_t reordered_filter_size = nchwc_output_channels * filter_input_channels; for (size_t i = 2; i < 4; i++) { diff --git a/onnxruntime/core/session/abi_session_options.cc b/onnxruntime/core/session/abi_session_options.cc index cb721b65f8..d8aadcdb7d 100644 --- a/onnxruntime/core/session/abi_session_options.cc +++ b/onnxruntime/core/session/abi_session_options.cc @@ -3,6 +3,7 @@ #include "core/graph/onnx_protobuf.h" #include "core/common/gsl_suppress.h" +#include "core/common/inlined_containers.h" #include "core/session/onnxruntime_c_api.h" #include "core/session/ort_apis.h" #include "core/framework/error_code_helper.h" @@ -177,9 +178,47 @@ ORT_API_STATUS_IMPL(OrtApis::AddSessionConfigEntry, _Inout_ OrtSessionOptions* o ORT_API_STATUS_IMPL(OrtApis::AddInitializer, _Inout_ OrtSessionOptions* options, _In_z_ const char* name, _In_ const OrtValue* val) { + API_IMPL_BEGIN auto st = options->value.AddInitializer(name, val); if (!st.IsOK()) { return onnxruntime::ToOrtStatus(st); } return nullptr; -} \ No newline at end of file + API_IMPL_END +} + +ORT_API_STATUS_IMPL(OrtApis::AddExternalInitializers, _In_ OrtSessionOptions* options, + _In_reads_(input_len) const char* const* initializer_names, + _In_reads_(input_len) const OrtValue* const* initializers, size_t initializers_num) { + +#if !defined(ORT_MINIMAL_BUILD) && !defined(DISABLE_EXTERNAL_INITIALIZERS) + API_IMPL_BEGIN + onnxruntime::InlinedVector names; + onnxruntime::InlinedVector values; + names.reserve(initializers_num); + values.reserve(initializers_num); + for (size_t i = 0; i < initializers_num; ++i) { + if (initializer_names[i] == nullptr || initializers[i] == nullptr) { + auto message = onnxruntime::MakeString("Input index: ", i, " contains null pointers"); + return OrtApis::CreateStatus(ORT_INVALID_ARGUMENT, message.c_str()); + } + names.emplace_back(initializer_names[i]); + values.emplace_back(*initializers[i]); + } + + auto st = options->value.AddExternalInitializers(names, values); + if (!st.IsOK()) { + return onnxruntime::ToOrtStatus(st); + } + return nullptr; + API_IMPL_END +#else + ORT_UNUSED_PARAMETER(options); + ORT_UNUSED_PARAMETER(initializer_names); + ORT_UNUSED_PARAMETER(initializers); + ORT_UNUSED_PARAMETER(initializers_num); + return OrtApis::CreateStatus(ORT_NOT_IMPLEMENTED, "External initializers are not supported in this build"); +#endif +} + + diff --git a/onnxruntime/core/session/inference_session.cc b/onnxruntime/core/session/inference_session.cc index aec976888e..24870d9403 100644 --- a/onnxruntime/core/session/inference_session.cc +++ b/onnxruntime/core/session/inference_session.cc @@ -1237,6 +1237,13 @@ common::Status InferenceSession::Initialize() { // re-acquire mutex std::lock_guard l(session_mutex_); +#if !defined(DISABLE_EXTERNAL_INITIALIZERS) && !defined(ORT_MINIMAL_BUILD) + if (!session_options_.external_initializers.empty()) { + ORT_RETURN_IF_ERROR_SESSIONID_(graph.InjectExternalInitializedTensors(session_options_.external_initializers)); + InlinedHashMap{}.swap(session_options_.external_initializers); + } +#endif + // At this time we know all the providers that will be part of this session. // Read shared allocators from the environment and update them in the respective providers. // diff --git a/onnxruntime/core/session/onnxruntime_c_api.cc b/onnxruntime/core/session/onnxruntime_c_api.cc index 2a3800488a..fe8580f1a0 100644 --- a/onnxruntime/core/session/onnxruntime_c_api.cc +++ b/onnxruntime/core/session/onnxruntime_c_api.cc @@ -2521,6 +2521,7 @@ static constexpr OrtApi ort_api_1_to_12 = { &OrtApis::ReleaseCUDAProviderOptions, &OrtApis::SessionOptionsAppendExecutionProvider_MIGraphX, // End of Version 11 - DO NOT MODIFY ABOVE (see above text for more information) + &OrtApis::AddExternalInitializers, }; // Asserts to do a some checks to ensure older Versions of the OrtApi never change (will detect an addition or deletion but not if they cancel out each other) diff --git a/onnxruntime/core/session/ort_apis.h b/onnxruntime/core/session/ort_apis.h index c93f2cd33f..511255c81c 100644 --- a/onnxruntime/core/session/ort_apis.h +++ b/onnxruntime/core/session/ort_apis.h @@ -339,4 +339,8 @@ ORT_API_STATUS_IMPL(UpdateCUDAProviderOptions, _Inout_ OrtCUDAProviderOptionsV2* size_t num_keys); ORT_API_STATUS_IMPL(GetCUDAProviderOptionsAsString, _In_ const OrtCUDAProviderOptionsV2* cuda_options, _Inout_ OrtAllocator* allocator, _Outptr_ char** ptr); ORT_API(void, ReleaseCUDAProviderOptions, _Frees_ptr_opt_ OrtCUDAProviderOptionsV2*); +ORT_API_STATUS_IMPL(AddExternalInitializers, _In_ OrtSessionOptions* options, + _In_reads_(input_len) const char* const* input_names, + _In_reads_(input_len) const OrtValue* const* inputs, size_t input_len); + } // namespace OrtApis diff --git a/onnxruntime/python/onnxruntime_pybind_state.cc b/onnxruntime/python/onnxruntime_pybind_state.cc index 5b11bf1bc5..b255b26552 100644 --- a/onnxruntime/python/onnxruntime_pybind_state.cc +++ b/onnxruntime/python/onnxruntime_pybind_state.cc @@ -9,6 +9,7 @@ #define PY_ARRAY_UNIQUE_SYMBOL onnxruntime_python_ARRAY_API #include +#include "core/common/inlined_containers.h" #include "core/common/logging/logging.h" #include "core/common/logging/severity.h" #include "core/common/optional.h" @@ -1234,7 +1235,28 @@ Applies to session load, initialization, etc. Default is 0.)pbdoc") // This is no different than the native APIs const OrtValue* ml_value = ml_value_pyobject.attr(PYTHON_ORTVALUE_NATIVE_OBJECT_ATTR).cast(); ORT_THROW_IF_ERROR(options->AddInitializer(name, ml_value)); - }); + }) + .def("add_external_initializers", [](PySessionOptions* options, py::list& names, + const py::list& ort_values) -> void { +#if !defined(ORT_MINIMAL_BUILD) && !defined(DISABLE_EXTERNAL_INITIALIZERS) + const auto init_num = ort_values.size(); + ORT_ENFORCE(init_num == names.size(), "Expecting names and ort_values lists to have equal length"); + InlinedVector names_ptrs; + InlinedVector values_ptrs; + names_ptrs.reserve(init_num); + values_ptrs.reserve(init_num); + for (size_t i = 0; i < init_num; ++i) { + names_ptrs.emplace_back(py::str(names[i])); + values_ptrs.emplace_back(*ort_values[i].attr(PYTHON_ORTVALUE_NATIVE_OBJECT_ATTR).cast()); + } + ORT_THROW_IF_ERROR(options->AddExternalInitializers(names_ptrs, values_ptrs)); +#else + ORT_UNUSED_PARAMETER(options); + ORT_UNUSED_PARAMETER(names); + ORT_UNUSED_PARAMETER(ort_values); + ORT_THROW("External initializers are not supported in this build."); +#endif + }); py::class_(m, "RunOptions", R"pbdoc(Configuration information for a single Run.)pbdoc") .def(py::init()) diff --git a/onnxruntime/test/ir/graph_test.cc b/onnxruntime/test/ir/graph_test.cc index 96b725c755..36d6fb9b6b 100644 --- a/onnxruntime/test/ir/graph_test.cc +++ b/onnxruntime/test/ir/graph_test.cc @@ -2,6 +2,8 @@ // Licensed under the MIT License. #include +#include "core/common/inlined_containers.h" +#include "core/framework/tensorprotoutils.h" #include "core/graph/graph_viewer.h" #include "core/graph/model.h" #include "core/graph/op.h" @@ -1697,7 +1699,7 @@ TEST_F(GraphTest, ReplaceInitializedTensor) { ONNX_NAMESPACE::TensorProto bad_name = original; bad_name.set_name("invalid"); - status = graph.ReplaceInitializedTensor(bad_name); + status = graph.ReplaceInitializedTensor(std::move(bad_name)); ASSERT_FALSE(status.IsOK()); } @@ -1705,7 +1707,7 @@ TEST_F(GraphTest, ReplaceInitializedTensor) { ONNX_NAMESPACE::TensorProto bad_type = original; bad_type.set_data_type(TensorProto_DataType_FLOAT16); - status = graph.ReplaceInitializedTensor(bad_type); + status = graph.ReplaceInitializedTensor(std::move(bad_type)); ASSERT_FALSE(status.IsOK()); } @@ -1715,7 +1717,7 @@ TEST_F(GraphTest, ReplaceInitializedTensor) { bad_dims.add_dims(2); bad_dims.add_dims(1); - status = graph.ReplaceInitializedTensor(bad_dims); + status = graph.ReplaceInitializedTensor(std::move(bad_dims)); ASSERT_FALSE(status.IsOK()); } @@ -1748,6 +1750,88 @@ TEST_F(GraphTest, ReplaceInitializedTensor) { } } +#if !defined(ORT_MINIMAL_BUILD) && !defined(DISABLE_EXTERNAL_INITIALIZERS) + +namespace { +void SetTensorProtoExternalData(const std::string& key, const std::string& value, + ONNX_NAMESPACE::TensorProto& tensor_proto) { + auto* external_data = tensor_proto.mutable_external_data(); + auto kvp_it = std::find_if( + external_data->begin(), external_data->end(), + [&key](const ONNX_NAMESPACE::StringStringEntryProto& kvp) { return kvp.key() == key; }); + auto* kvp = kvp_it != external_data->end() ? &(*kvp_it) : external_data->Add(); + kvp->set_key(key); + kvp->set_value(value); +} +} // namespace + +TEST_F(GraphTest, InjectExternalInitializedTensors) { + const std::string initializer_name = "test_external_initializer"; + + std::vector tensor_data = []() { + std::vector tensor_data(100); + std::iota(tensor_data.begin(), tensor_data.end(), 0); + return tensor_data; + }(); + + // Create OrtValue for replacement + TensorShape data_shape{static_cast(tensor_data.size())}; + OrtValue ort_value; + Tensor::InitOrtValue(DataTypeImpl::GetType(), data_shape, tensor_data.data(), + OrtMemoryInfo(onnxruntime::CPU, OrtAllocatorType::OrtDeviceAllocator), ort_value); + const InlinedHashMap injection_initializers = { + {initializer_name, ort_value} + }; + + // We do not need actual files there since we are not going to load it. + const auto tensor_data_dir_path = Path::Parse(ToPathString(".")); + const auto tensor_data_dir_relative_path = Path::Parse(ToPathString("external_data.bin")); + + const auto tensor_proto = + [&]() { + ONNX_NAMESPACE::TensorProto tensor_proto; + tensor_proto.set_name(initializer_name); + tensor_proto.add_dims(tensor_data.size()); + tensor_proto.set_data_type(ONNX_NAMESPACE::TensorProto_DataType_INT32); + tensor_proto.set_data_location(ONNX_NAMESPACE::TensorProto_DataLocation_EXTERNAL); + SetTensorProtoExternalData("location", ToUTF8String(tensor_data_dir_relative_path.ToPathString()), tensor_proto); + SetTensorProtoExternalData("offset", "0", tensor_proto); + SetTensorProtoExternalData("length", std::to_string(tensor_data.size() * sizeof(int32_t)), tensor_proto); + return tensor_proto; + }(); + + Model m{"test_model", false, *logger_}; + Graph& graph = m.MainGraph(); + graph.AddInitializedTensor(tensor_proto); + + ASSERT_EQ(graph.GetAllInitializedTensors().size(), 1U); + + const TensorProto* external_data = nullptr; + ASSERT_TRUE(graph.GetInitializedTensor(initializer_name, external_data)); + ASSERT_TRUE(utils::HasExternalData(*external_data)); + + // Replace things. + ASSERT_STATUS_OK(graph.InjectExternalInitializedTensors(injection_initializers)); + + ASSERT_EQ(graph.GetAllInitializedTensors().size(), 1U); + + const TensorProto* with_data = nullptr; + ASSERT_TRUE(graph.GetInitializedTensor(initializer_name, with_data)); + // No longer has external data + ASSERT_FALSE(utils::HasExternalData(*with_data)); + + const auto& original_tensor = ort_value.Get(); + + Tensor replaced_tensor(original_tensor.DataType(), data_shape, std::make_shared()); + ASSERT_STATUS_OK(utils::TensorProtoToTensor(Env::Default(), tensor_data_dir_path.ToPathString().c_str(), *with_data, replaced_tensor)); + + ASSERT_EQ(original_tensor.GetElementType(), replaced_tensor.GetElementType()); + const auto original_span = original_tensor.DataAsSpan(); + const auto replaced_span = replaced_tensor.DataAsSpan(); + ASSERT_EQ(original_span, replaced_span); +} +#endif + TEST_F(GraphTest, AddRemoveInitializerHandling) { Model m{"test_model", false, *logger_}; Graph& graph = m.MainGraph(); diff --git a/onnxruntime/test/python/onnxruntime_test_python.py b/onnxruntime/test/python/onnxruntime_test_python.py index 858f18c97e..8f29c974f9 100644 --- a/onnxruntime/test/python/onnxruntime_test_python.py +++ b/onnxruntime/test/python/onnxruntime_test_python.py @@ -827,6 +827,16 @@ class TestInferenceSession(unittest.TestCase): res = sess.run(["Y"], {"X": np.array([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]], dtype=np.float32)}) self.assertTrue(np.array_equal(res[0], np.array([[2.0, 2.0], [12.0, 12.0], [30.0, 30.0]], dtype=np.float32))) + def testSessionOptionsAddExternalInitializers(self): + # Create an external initializer data in OrtValue + # This initializer will replace the initializer with external data reference in the graph + ortvalue_initializer = onnxrt.OrtValue.ortvalue_from_numpy(np.array([0, 0, 1, 1]).astype(np.int64)) + so = onnxrt.SessionOptions() + so.add_external_initializers(["Pads_not_on_disk"], [ortvalue_initializer]) + # This should not throw + onnxrt.InferenceSession(get_name("model_with_external_initializer_come_from_user.onnx"), sess_options=so, providers=['CPUExecutionProvider']) + + def testRegisterCustomOpsLibrary(self): if sys.platform.startswith("win"): shared_library = 'custom_op_library.dll' diff --git a/onnxruntime/test/shared_lib/test_model_loading.cc b/onnxruntime/test/shared_lib/test_model_loading.cc index 9457cf9029..5324b1ffb8 100644 --- a/onnxruntime/test/shared_lib/test_model_loading.cc +++ b/onnxruntime/test/shared_lib/test_model_loading.cc @@ -60,7 +60,7 @@ TEST(CApiTest, model_from_array) { #ifdef DISABLE_EXTERNAL_INITIALIZERS TEST(CApiTest, TestDisableExternalInitiliazers) { - const char* model_path = "testdata/model_with_external_initializers.onnx"; + constexpr auto model_path = ORT_TSTR("testdata/model_with_external_initializers.onnx"); Ort::SessionOptions so; try { @@ -70,6 +70,25 @@ TEST(CApiTest, TestDisableExternalInitiliazers) { ASSERT_THAT(ex.what(), testing::HasSubstr("Initializer tensors with external data is not allowed.")); } } + +#elif !defined(ORT_MINIMAL_BUILD) +TEST(CApiTest, TestExternalInitializersInjection) { + constexpr auto model_path = ORT_TSTR("testdata/model_with_external_initializer_come_from_user.onnx"); + std::array Pads_not_on_disk{0, 0, 1, 1}; + constexpr std::array init_shape{4}; + + const std::vector init_names{"Pads_not_on_disk"}; + std::vector initializer_data; + + auto cpu_mem_info = Ort::MemoryInfo::CreateCpu(OrtDeviceAllocator, OrtMemTypeCPU); + auto init_tensor = Ort::Value::CreateTensor(cpu_mem_info, Pads_not_on_disk.data(), Pads_not_on_disk.size(), init_shape.data(), init_shape.size()); + initializer_data.push_back(std::move(init_tensor)); + + Ort::SessionOptions so; + so.AddExternalInitializers(init_names, initializer_data); + EXPECT_NO_THROW(Ort::Session(*ort_env, model_path, so)); +} + #endif } // namespace test } // namespace onnxruntime diff --git a/onnxruntime/test/testdata/model_with_external_initializer_come_from_user.onnx b/onnxruntime/test/testdata/model_with_external_initializer_come_from_user.onnx new file mode 100644 index 0000000000..486573918e --- /dev/null +++ b/onnxruntime/test/testdata/model_with_external_initializer_come_from_user.onnx @@ -0,0 +1,19 @@ + onnx-example:Æ +2 +X +Pads_not_on_diskY"Pad* +mode"constant  +test-model*:BPads_not_on_diskj +locationPads_not_on_disk.binpZ +X +  + +Z +Pads_not_on_disk + + +b +Y +  + +B \ No newline at end of file diff --git a/onnxruntime/test/testdata/model_with_external_initializer_come_from_user.py b/onnxruntime/test/testdata/model_with_external_initializer_come_from_user.py new file mode 100644 index 0000000000..baad806b7e --- /dev/null +++ b/onnxruntime/test/testdata/model_with_external_initializer_come_from_user.py @@ -0,0 +1,62 @@ +# Copyright (c) Microsoft Corporation. All rights reserved. +# Licensed under the MIT License. + +import onnx +import numpy as np +from onnx import helper +from onnx import TensorProto +from onnx.numpy_helper import from_array +from onnx.external_data_helper import set_external_data + + +def create_external_data_tensor(value, tensor_name): # type: (List[Any], Text) -> TensorProto + tensor = from_array(value) + tensor.name = tensor_name + tensor_filename = "{}.bin".format(tensor_name) + set_external_data(tensor, location=tensor_filename) + tensor.ClearField('raw_data') + tensor.data_location = onnx.TensorProto.EXTERNAL + return tensor + + +def GenerateModel(model_name): + # Create one input (ValueInfoProto) + X = helper.make_tensor_value_info('X', TensorProto.FLOAT, [1, 2]) + + # Create second input (ValueInfoProto) + Pads = helper.make_tensor_value_info('Pads_not_on_disk', TensorProto.INT64, [4]) + + # Create one output (ValueInfoProto) + Y = helper.make_tensor_value_info('Y', TensorProto.FLOAT, [1, 4]) + + # Create a node (NodeProto) + node_def = helper.make_node( + 'Pad', # node name + ['X', 'Pads_not_on_disk'], # inputs + ['Y'], # outputs + mode='constant', # Attributes + ) + + # Create the graph (GraphProto) + initializer_data = np.array([0, 0, 1, 1]).astype(np.int64) + graph_def = helper.make_graph( + [node_def], + "test-model", + [X, Pads], + [Y], + [create_external_data_tensor(initializer_data, "Pads_not_on_disk")] + ) + + # Create the model (ModelProto) + model_def = helper.make_model(graph_def, + producer_name='onnx-example') + + print('The ir_version in model: {}\n'.format(model_def.ir_version)) + print('The producer_name in model: {}\n'.format(model_def.producer_name)) + print('The graph in model:\n{}'.format(model_def.graph)) + with open(model_name, "wb") as model_file: + model_file.write(model_def.SerializeToString()) + + +if __name__ == "__main__": + GenerateModel('model_with_external_initializer_come_from_user.onnx') diff --git a/orttraining/orttraining/models/bert/main.cc b/orttraining/orttraining/models/bert/main.cc index 2c303086c0..18004ef4bc 100644 --- a/orttraining/orttraining/models/bert/main.cc +++ b/orttraining/orttraining/models/bert/main.cc @@ -58,7 +58,8 @@ static SessionOptions session_options = { true, //thread_pool_allow_spinning false, //use_deterministic_compute {}, //config_options - {}, // initializers_to_share_map + {}, //initializers_to_share_map + {}, // external_initializers }; struct BertParameters : public TrainingRunner::Parameters { diff --git a/orttraining/orttraining/models/pipeline_poc/main.cc b/orttraining/orttraining/models/pipeline_poc/main.cc index 0e4fe1716d..25f7180be2 100644 --- a/orttraining/orttraining/models/pipeline_poc/main.cc +++ b/orttraining/orttraining/models/pipeline_poc/main.cc @@ -108,7 +108,8 @@ int main(int argc, char* argv[]) { true, //thread_pool_allow_spinning false, //use_deterministic_compute {}, //session_configurations - {} //initializers_to_share_map + {}, //initializers_to_share_map + {} // external_initializers }; InferenceSession session_object{so, *env}; diff --git a/orttraining/orttraining/models/runner/training_runner.cc b/orttraining/orttraining/models/runner/training_runner.cc index 664b489c62..b97ffc6f41 100644 --- a/orttraining/orttraining/models/runner/training_runner.cc +++ b/orttraining/orttraining/models/runner/training_runner.cc @@ -55,6 +55,7 @@ static SessionOptions SESSION_OPTION = { false, //use_deterministic_compute {}, //config_options {}, // initializers_to_share_map + {}, // external_initializers }; TrainingRunner::TrainingRunner(Parameters params, const Environment& env)