Fix UTs for SaveWithExternalInitializers

microsoft · Jan 4, 2024 · 794fd3b · 794fd3b
1 parent 81d6cbb
commit 794fd3b
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 8 deletions.
diff --git a/onnxruntime/core/graph/model.cc b/onnxruntime/core/graph/model.cc
@@ -81,8 +81,8 @@ Model::Model(const std::string& graph_name,
              const std::vector<ONNX_NAMESPACE::FunctionProto>& model_local_functions,
              const logging::Logger& logger,
              const ModelOptions& options)
-    : model_path_(Path::Parse(model_path)) {
-  external_data_path_ = model_path_.ParentPath();
+    : model_path_(std::move(Path::Parse(model_path)))
+    , external_data_path_(std::move(model_path_.ParentPath())) {
   model_proto_.set_ir_version(ONNX_NAMESPACE::Version::IR_VERSION);
   model_proto_.mutable_graph()->set_name(graph_name);
   model_metadata_ = model_metadata;
@@ -160,8 +160,8 @@ Model::Model(const ModelProto& model_proto, const PathString& model_path,
 Model::Model(ModelProto&& model_proto, const PathString& model_path,
              const IOnnxRuntimeOpSchemaRegistryList* local_registries,
              const logging::Logger& logger, const ModelOptions& options)
-    : model_path_(Path::Parse(model_path)) {
-  external_data_path_ = model_path_.ParentPath();
+    : model_path_(std::move(Path::Parse(model_path)))
+    , external_data_path_(std::move(model_path_.ParentPath())) {
   if (!utils::HasGraph(model_proto)) {
     ORT_THROW("ModelProto does not have a graph.");
   }

diff --git a/onnxruntime/test/framework/save_model_with_external_initializers.cc b/onnxruntime/test/framework/save_model_with_external_initializers.cc
@@ -44,22 +44,24 @@ void LoadSaveAndCompareModel(const std::string& input_onnx,
   // Compare the initializers of the two versions.
   Path model_path{};
   Path external_data_path{};
+  Path external_data_folder_path{};
   for (auto i : initializers) {
     const std::string kInitName = i.first;
     const ONNX_NAMESPACE::TensorProto* tensor_proto = i.second;
     const ONNX_NAMESPACE::TensorProto* from_external_tensor_proto = initializers_from_external[kInitName];
 
     std::vector<uint8_t> tensor_proto_data;
     model_path = Path::Parse(ToPathString(input_onnx));
-    external_data_path = model_path.ParentPath();
+    external_data_folder_path = model_path.ParentPath();
     // the file name of the external initializer data is tracked inside the tensor protocol
-    ORT_THROW_IF_ERROR(utils::UnpackInitializerData(*tensor_proto, external_data_path, tensor_proto_data));
+    ORT_THROW_IF_ERROR(utils::UnpackInitializerData(*tensor_proto, external_data_folder_path, tensor_proto_data));
     size_t tensor_proto_size = tensor_proto_data.size();
 
     std::vector<uint8_t> from_external_tensor_proto_data;
     model_path = Path::Parse(ToPathString(output_onnx));
-    external_data_path = model_path.ParentPath();
-    ORT_THROW_IF_ERROR(utils::UnpackInitializerData(*from_external_tensor_proto, external_data_path, from_external_tensor_proto_data));
+    external_data_folder_path = model_path.ParentPath();
+    external_data_path = model_path.ParentPath().Append(Path::Parse(ToPathString(output_external_init_file)));
+    ORT_THROW_IF_ERROR(utils::UnpackInitializerData(*from_external_tensor_proto, external_data_folder_path, from_external_tensor_proto_data));
     size_t from_external_tensor_proto_size = from_external_tensor_proto_data.size();
 
     if (from_external_tensor_proto_size < initializer_size_threshold) {