mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-06-08 00:23:03 +00:00
5daeb5e0b0
### Description Background: User save large model with initializer data in external file. e.g: onnx.save_model(onnx_model, "path/to/save/the/model.onnx", save_as_external_data=True, all_tensors_to_one_file=True, location="filename", size_threshold=1024). In that case, Ort loads the model, get the external initializer information (external file name, offset, length) and use the model path to find the external file, and locate to the tensor data via the offset and length. But it won't work if user load the model from memory, since Ort lost track of the model path. This PR adds API/session option to let user provide a table with external initializer file name as the key, the pointer to the loaded external file in memory and the buffer length as value. So that 1. user can load the model from memory buffer with external initializers in memory buffer too. 2. the initializers can be shared across sessions, for different EPs. 3. user can load the file in any way they want, e.g mmap. Internally, 1. at session creation time, Ort goes through the external initializers in the graph, gets the file name, offset, data length of the external initializers from Tensorproto . 2. With the file name, Ort get the file in memory buffer and buffer length from the table user provided. 4. Ort locates the tensor buffer from file in memory buffer (user provided) using the offset and data length (from Tensorproto ). 5. Ort creates the Tensor and replace the existing Tensor in the graph. ### Motivation and Context https://github.com/onnx/onnx/blob/main/docs/ExternalData.md For a model with external data, the Tensorproto may have initializer data in a separate file. The external file location is set via the file path relative to the model path. With the API to load model from memory buffer, it lost track of the model path. So it causes error if the model has external data. By adding a session option to set the external data buffer, Ort can find the external data correctly if model loaded from memory buffer.
99 lines
3.8 KiB
C++
99 lines
3.8 KiB
C++
// Copyright (c) Microsoft Corporation. All rights reserved.
|
|
// Licensed under the MIT License.
|
|
|
|
#include "core/framework/session_options.h"
|
|
#include "core/common/logging/logging.h"
|
|
#include "core/framework/ort_value.h"
|
|
|
|
namespace onnxruntime {
|
|
|
|
namespace {
|
|
|
|
Status CheckInitializer(const char* name, const OrtValue* val) {
|
|
if (name == nullptr) {
|
|
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");
|
|
}
|
|
|
|
if (!val->IsTensor()) {
|
|
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "Received OrtValue is not a tensor. Only tensors are supported.");
|
|
}
|
|
if (val->Get<Tensor>().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
|
|
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<const std::string> names, gsl::span<const OrtValue> 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();
|
|
}
|
|
|
|
Status SessionOptions::AddExternalInitializersFromFilesInMemory(gsl::span<const PathString> file_names,
|
|
gsl::span<std::pair<char*, const size_t>> files_buffers) {
|
|
const auto num_files = file_names.size();
|
|
ORT_ENFORCE(num_files == files_buffers.size(), "Expecting same size spans");
|
|
external_initializer_files_mmap.reserve(external_initializer_files_mmap.size() + num_files);
|
|
static constexpr std::array<std::basic_string_view<ORTCHAR_T>, 4> prefix_list{
|
|
ORT_TSTR(".//"),
|
|
ORT_TSTR("./"),
|
|
ORT_TSTR(".\\\\"),
|
|
ORT_TSTR(".\\")};
|
|
for (size_t i = 0; i < num_files; ++i) {
|
|
// ignore "./" from file name if it has
|
|
auto file_name = file_names[i];
|
|
for (auto prefix : prefix_list) {
|
|
if (file_name.rfind(prefix, 0) == 0) {
|
|
file_name = file_name.substr(prefix.length());
|
|
break;
|
|
}
|
|
}
|
|
bool result = external_initializer_files_mmap.emplace(file_name, files_buffers[i]).second;
|
|
if (!result) {
|
|
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "An entry for this name has already been added: ",
|
|
ORT_TSTR_CONVERT_TO_PRINTABLE_STRING(file_name));
|
|
}
|
|
}
|
|
return Status::OK();
|
|
}
|
|
#endif // !defined(ORT_MINIMAL_BUILD) && !defined(DISABLE_EXTERNAL_INITIALIZERS)
|
|
|
|
#if !defined(ORT_MINIMAL_BUILD) || defined(ORT_MINIMAL_BUILD_CUSTOM_OPS)
|
|
void SessionOptions::AddCustomOpLibraryHandle(PathString library_name, void* library_handle) {
|
|
if (!this->custom_op_libs) {
|
|
this->custom_op_libs = std::make_shared<LibraryHandles>();
|
|
}
|
|
|
|
this->custom_op_libs->Add(std::move(library_name), library_handle);
|
|
}
|
|
#endif // !defined(ORT_MINIMAL_BUILD) || defined(ORT_MINIMAL_BUILD_CUSTOM_OPS)
|
|
|
|
} // namespace onnxruntime
|