Skip to content

Commit

Permalink
Enable user loading model with external data from memory buffer (#23557)
Browse files Browse the repository at this point in the history
Add session option to enable user loading model with external data from memory buffer. User want to set the folder path for the external data files.

### Description
For some cases user load the model from memory buffer, but they can't load the external files into memory. They need to have a way to set the folder path for the external data files so that Ort can figure out the external data location.
  • Loading branch information
HectorSVC authored Feb 6, 2025
1 parent 544bdd6 commit 649ced4
Show file tree
Hide file tree
Showing 3 changed files with 76 additions and 3 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -250,6 +250,12 @@ static const char* const kOrtSessionOptionsOptimizedModelExternalInitializersFil
static const char* const kOrtSessionOptionsOptimizedModelExternalInitializersMinSizeInBytes =
"session.optimized_model_external_initializers_min_size_in_bytes";

// When loading model from memory buffer and the model has external initializers
// Use this config to set the external data file folder path
// All external data files should be in the same folder
static const char* const kOrtSessionOptionsModelExternalInitializersFileFolderPath =
"session.model_external_initializers_file_folder_path";

// Use this config when saving pre-packed constant initializers to an external data file.
// This allows you to memory map pre-packed initializers on model load and leave it to
// to the OS the amount of memory consumed by the pre-packed initializers. Otherwise,
Expand Down
27 changes: 24 additions & 3 deletions onnxruntime/core/session/inference_session.cc
Original file line number Diff line number Diff line change
Expand Up @@ -1090,7 +1090,14 @@ common::Status InferenceSession::Load(const void* model_data, int model_data_len

const bool strict_shape_type_inference = session_options_.config_options.GetConfigOrDefault(
kOrtSessionOptionsConfigStrictShapeTypeInference, "0") == "1";
return onnxruntime::Model::Load(std::move(model_proto), PathString(), model,

std::string external_data_folder_path = session_options_.config_options.GetConfigOrDefault(
kOrtSessionOptionsModelExternalInitializersFileFolderPath, "");
if (!external_data_folder_path.empty() && model_location_.empty()) {
model_location_ = ToPathString(external_data_folder_path + "/virtual_model.onnx");
}

return onnxruntime::Model::Load(std::move(model_proto), model_location_, model,
HasLocalSchema() ? &custom_schema_registries_ : nullptr, *session_logger_,
ModelOptions(true, strict_shape_type_inference));
};
Expand Down Expand Up @@ -1120,8 +1127,15 @@ common::Status InferenceSession::LoadOnnxModel(ModelProto model_proto) {
#endif
const bool strict_shape_type_inference = session_options_.config_options.GetConfigOrDefault(
kOrtSessionOptionsConfigStrictShapeTypeInference, "0") == "1";

std::string external_data_folder_path = session_options_.config_options.GetConfigOrDefault(
kOrtSessionOptionsModelExternalInitializersFileFolderPath, "");
if (!external_data_folder_path.empty() && model_location_.empty()) {
model_location_ = ToPathString(external_data_folder_path + "/virtual_model.onnx");
}

// This call will move model_proto to the constructed model instance
return onnxruntime::Model::Load(std::move(model_proto), PathString(), model,
return onnxruntime::Model::Load(std::move(model_proto), model_location_, model,
HasLocalSchema() ? &custom_schema_registries_ : nullptr, *session_logger_,
ModelOptions(true, strict_shape_type_inference));
};
Expand Down Expand Up @@ -1157,7 +1171,14 @@ common::Status InferenceSession::Load(std::istream& model_istream, bool allow_re
kOrtSessionOptionsConfigStrictShapeTypeInference, "0") == "1";
ModelOptions model_opts(allow_released_opsets_only,
strict_shape_type_inference);
return onnxruntime::Model::Load(std::move(model_proto), PathString(), model,

std::string external_data_folder_path = session_options_.config_options.GetConfigOrDefault(
kOrtSessionOptionsModelExternalInitializersFileFolderPath, "");
if (!external_data_folder_path.empty() && model_location_.empty()) {
model_location_ = ToPathString(external_data_folder_path + "/virtual_model.onnx");
}

return onnxruntime::Model::Load(std::move(model_proto), model_location_, model,
HasLocalSchema() ? &custom_schema_registries_ : nullptr,
*session_logger_, model_opts);
};
Expand Down
46 changes: 46 additions & 0 deletions onnxruntime/test/shared_lib/test_model_loading.cc
Original file line number Diff line number Diff line change
Expand Up @@ -215,6 +215,52 @@ TEST(CApiTest, TestLoadModelFromArrayWithExternalInitializersFromFileArrayPathRo
#endif
}

// The model has external data, Test loading model from array
// Extra API required to set the external data path
TEST(CApiTest, TestLoadModelFromArrayWithExternalInitializersViaSetExternalDataPath) {
std::string model_file_name = "conv_qdq_external_ini.onnx";
std::string external_bin_name = "conv_qdq_external_ini.bin";
std::string test_folder = "testdata/";
std::string model_path = test_folder + model_file_name;
std::vector<char> buffer;
ReadFileToBuffer(model_path.c_str(), buffer);

std::vector<char> external_bin_buffer;
std::string external_bin_path = test_folder + external_bin_name;
ReadFileToBuffer(external_bin_path.c_str(), external_bin_buffer);

Ort::SessionOptions so;
std::string optimized_model_file_name(model_file_name);
auto length = optimized_model_file_name.length();
optimized_model_file_name.insert(length - 5, "_opt");
std::string optimized_file_path(test_folder + optimized_model_file_name);
PathString optimized_file_path_t(optimized_file_path.begin(), optimized_file_path.end());

// Dump the optimized model with external data so that it will unpack the external data from the loaded model
so.SetOptimizedModelFilePath(optimized_file_path_t.c_str());

// set the model external file folder path
so.AddConfigEntry(kOrtSessionOptionsModelExternalInitializersFileFolderPath, test_folder.c_str());

std::string opt_bin_file_name(optimized_model_file_name);
opt_bin_file_name.replace(optimized_model_file_name.length() - 4, 4, "bin");
so.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_DISABLE_ALL);
so.AddConfigEntry(kOrtSessionOptionsOptimizedModelExternalInitializersFileName, opt_bin_file_name.c_str());
so.AddConfigEntry(kOrtSessionOptionsOptimizedModelExternalInitializersMinSizeInBytes, "10");

Ort::Session session(*ort_env.get(), buffer.data(), buffer.size(), so);

std::string generated_bin_path = test_folder + opt_bin_file_name;
std::vector<char> generated_bin_buffer;
ReadFileToBuffer(generated_bin_path.c_str(), generated_bin_buffer);

ASSERT_EQ(external_bin_buffer, generated_bin_buffer);

// Cleanup.
ASSERT_EQ(std::remove(optimized_file_path.c_str()), 0);
ASSERT_EQ(std::remove(generated_bin_path.c_str()), 0);
}

#ifndef _WIN32
struct FileDescriptorTraits {
using Handle = int;
Expand Down

0 comments on commit 649ced4

Please sign in to comment.