diff --git a/include/onnxruntime/core/graph/graph.h b/include/onnxruntime/core/graph/graph.h index be92550efe..d1841ff43b 100644 --- a/include/onnxruntime/core/graph/graph.h +++ b/include/onnxruntime/core/graph/graph.h @@ -108,6 +108,9 @@ class Node { /** Gets the domain of the OperatorSet that specifies the operator returned by #OpType. */ const std::string& Domain() const noexcept { return domain_; } + /** Gets the path of the owning model if any. */ + const Path& ModelPath() const noexcept; + /** Gets the Node's execution priority. @remarks Lower value means higher priority */ int Priority() const noexcept { return priority_; }; @@ -149,6 +152,7 @@ class Node { /** Gets the function body if applicable otherwise nullptr. */ const Function* GetFunctionBody() const noexcept { return func_body_; } + #endif /** diff --git a/include/onnxruntime/core/graph/graph_viewer.h b/include/onnxruntime/core/graph/graph_viewer.h index 78b114e8dd..b2d9ee22cc 100644 --- a/include/onnxruntime/core/graph/graph_viewer.h +++ b/include/onnxruntime/core/graph/graph_viewer.h @@ -42,6 +42,9 @@ class GraphViewer { /** Gets the Graph description. */ const std::string& Description() const noexcept; + /** Gets the path of the owning model if any **/ + const Path& ModelPath() const noexcept { return graph_->ModelPath(); } + /** Gets a tensor created from an initializer. @param tensor_name The tensor name diff --git a/onnxruntime/core/framework/tensorprotoutils.cc b/onnxruntime/core/framework/tensorprotoutils.cc index 3605487ba5..8d396aa071 100644 --- a/onnxruntime/core/framework/tensorprotoutils.cc +++ b/onnxruntime/core/framework/tensorprotoutils.cc @@ -114,10 +114,120 @@ static Status UnpackTensorWithRawData(const void* raw_data, size_t raw_data_leng gsl::make_span(raw_data_bytes, raw_data_length), gsl::make_span(p_data, expected_size))); return Status::OK(); } + +static Status GetExternalDataInfo(const ONNX_NAMESPACE::TensorProto& tensor_proto, + const ORTCHAR_T* tensor_proto_dir, + std::basic_string& external_file_path, + onnxruntime::FileOffsetType& file_offset, + SafeInt& tensor_data_length) { + 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."); + + 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(); + } + + file_offset = external_data_info->GetOffset(); + + ORT_RETURN_IF_ERROR(onnxruntime::utils::GetSizeInBytesFromTensorProto<0>( + tensor_proto, &tensor_data_length)); + const size_t external_data_length = external_data_info->GetLength(); + + ORT_RETURN_IF_NOT( + external_data_length == 0 || + external_data_length == tensor_data_length, + "TensorProto external data size mismatch. ", + "Computed size: ", *&tensor_data_length, + ", external_data.length: ", external_data_length); + + return Status::OK(); +} + +// Read external data for tensor in unint8_t* form and return Status::OK() if the data is read successfully. +// Uses the tensor_proto_dir to construct the full path for external data. If tensor_proto_dir == nullptr +// then uses the current directory instead. +// This function does not unpack string_data of an initializer tensor +static Status ReadExternalDataForTensor(const ONNX_NAMESPACE::TensorProto& tensor_proto, + const ORTCHAR_T* tensor_proto_dir, + std::unique_ptr& unpacked_tensor, + SafeInt& tensor_data_length) { + std::basic_string external_file_path; + onnxruntime::FileOffsetType file_offset; + ORT_RETURN_IF_ERROR(GetExternalDataInfo( + tensor_proto, + tensor_proto_dir, + external_file_path, + file_offset, + tensor_data_length)); + + unpacked_tensor.reset(new uint8_t[*&tensor_data_length]); + ORT_RETURN_IF_ERROR(onnxruntime::Env::Default().ReadFileIntoBuffer( + external_file_path.c_str(), + file_offset, + tensor_data_length, + gsl::make_span(reinterpret_cast(unpacked_tensor.get()), tensor_data_length))); + + return Status::OK(); +} } // namespace namespace onnxruntime { namespace utils { +#if !defined(ORT_MINIMAL_BUILD) +#define DEFINE_UNPACK_EXTERNAL_TENSOR(T) \ + template <> \ + Status UnpackTensorWithExternalData(const ONNX_NAMESPACE::TensorProto& tensor, \ + const ORTCHAR_T* tensor_proto_dir, size_t expected_size, \ + /*out*/ T* p_data) { \ + ORT_RETURN_IF(nullptr == p_data); \ + \ + std::unique_ptr unpacked_tensor; \ + SafeInt tensor_byte_size = 0; \ + ORT_RETURN_IF_ERROR(ReadExternalDataForTensor( \ + tensor, \ + tensor_proto_dir, \ + unpacked_tensor, \ + tensor_byte_size)); \ + \ + size_t element_count = tensor_byte_size / sizeof(T); \ + ORT_RETURN_IF_NOT(expected_size == element_count, "Expected data size does not match the actual external data size."); \ + ORT_RETURN_IF_ERROR(onnxruntime::utils::ReadLittleEndian( \ + gsl::make_span(reinterpret_cast(unpacked_tensor.get()), tensor_byte_size), \ + gsl::make_span(p_data, expected_size))); \ + \ + return Status::OK(); \ + } + +DEFINE_UNPACK_EXTERNAL_TENSOR(float) +DEFINE_UNPACK_EXTERNAL_TENSOR(double) +DEFINE_UNPACK_EXTERNAL_TENSOR(uint8_t) +DEFINE_UNPACK_EXTERNAL_TENSOR(int8_t) +DEFINE_UNPACK_EXTERNAL_TENSOR(int16_t) +DEFINE_UNPACK_EXTERNAL_TENSOR(uint16_t) +DEFINE_UNPACK_EXTERNAL_TENSOR(int32_t) +DEFINE_UNPACK_EXTERNAL_TENSOR(int64_t) +DEFINE_UNPACK_EXTERNAL_TENSOR(uint64_t) +DEFINE_UNPACK_EXTERNAL_TENSOR(uint32_t) +DEFINE_UNPACK_EXTERNAL_TENSOR(bool) +DEFINE_UNPACK_EXTERNAL_TENSOR(MLFloat16) +DEFINE_UNPACK_EXTERNAL_TENSOR(BFloat16) + +template <> +Status UnpackTensorWithExternalData(const ONNX_NAMESPACE::TensorProto& /*tensor*/, + const ORTCHAR_T* /*tensor_proto_dir*/, size_t /*expected_size*/, + /*out*/ std::string* /*p_data*/) { + return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, + "External data type cannot be STRING."); +} +#endif //!defined(ORT_MINIMAL_BUILD) // This macro doesn't work for Float16/bool/string tensors #define DEFINE_UNPACK_TENSOR(T, Type, field_name, field_size) \ @@ -424,7 +534,9 @@ static void MoveOrtCallback(OrtCallback& from, OrtCallback& to) { #pragma warning(push) #pragma warning(disable : 6239) #endif -Status TensorProtoToMLValue(const Env& env, const ORTCHAR_T* tensor_proto_path, +// TODO: Change the current interface to take Path object for model path +// so that validating and manipulating path for reading external data becomes easy +Status TensorProtoToMLValue(const Env& env, const ORTCHAR_T* model_path, const ONNX_NAMESPACE::TensorProto& tensor_proto, const MemBuffer& m, OrtValue& value, OrtCallback& deleter) { const OrtMemoryInfo& allocator = m.GetAllocInfo(); @@ -432,28 +544,27 @@ Status TensorProtoToMLValue(const Env& env, const ORTCHAR_T* tensor_proto_path, deleter.f = nullptr; deleter.param = nullptr; void* raw_data = nullptr; - size_t raw_data_len = 0; + SafeInt raw_data_len = 0; const DataTypeImpl* const type = DataTypeImpl::TensorTypeFromONNXEnum(tensor_proto.data_type())->GetElementType(); AutoDelete deleter_for_file_data; void* tensor_data; { - if (tensor_proto.data_location() == TensorProto_DataLocation_EXTERNAL) { - if (ele_type == ONNX_TENSOR_ELEMENT_DATA_TYPE_STRING) - return Status(common::ONNXRUNTIME, common::INVALID_ARGUMENT, "string tensor can not have raw data"); - - std::unique_ptr external_data_info; - ORT_RETURN_IF_ERROR(ExternalDataInfo::Create(tensor_proto.external_data(), external_data_info)); - std::basic_string full_path; - if (tensor_proto_path != nullptr) { - ORT_RETURN_IF_ERROR(GetDirNameFromFilePath(tensor_proto_path, full_path)); - full_path = ConcatPathComponent(full_path, external_data_info->GetRelPath()); - } else { - full_path = external_data_info->GetRelPath(); + if (utils::HasExternalData(tensor_proto)) { + // Get the external data info + std::basic_string external_data_file_path; + FileOffsetType file_offset; + std::basic_string tensor_proto_dir; + if (model_path != nullptr) { + ORT_RETURN_IF_ERROR(GetDirNameFromFilePath(model_path, tensor_proto_dir)); } - raw_data_len = external_data_info->GetLength(); + ORT_RETURN_IF_ERROR(GetExternalDataInfo( + tensor_proto, + tensor_proto_dir.size() == 0 ? nullptr : tensor_proto_dir.c_str(), + external_data_file_path, file_offset, raw_data_len)); + // load the file ORT_RETURN_IF_ERROR(GetFileContent( - env, full_path.c_str(), external_data_info->GetOffset(), raw_data_len, + env, external_data_file_path.c_str(), file_offset, raw_data_len, raw_data, deleter_for_file_data.d)); } else if (utils::HasRawData(tensor_proto)) { if (ele_type == ONNX_TENSOR_ELEMENT_DATA_TYPE_STRING) @@ -605,6 +716,7 @@ ONNX_NAMESPACE::TensorProto TensorToTensorProto(const Tensor& tensor, const std: } common::Status ConstantNodeProtoToTensorProto(const ONNX_NAMESPACE::NodeProto& node, + const Path& model_path, ONNX_NAMESPACE::TensorProto& tensor) { const AttributeProto& constant_attribute = node.attribute(0); @@ -639,7 +751,7 @@ common::Status ConstantNodeProtoToTensorProto(const ONNX_NAMESPACE::NodeProto& n } case AttributeProto_AttributeType_SPARSE_TENSOR: { auto& s = constant_attribute.sparse_tensor(); - ORT_RETURN_IF_ERROR(SparseTensorProtoToDenseTensorProto(s, tensor)); + ORT_RETURN_IF_ERROR(SparseTensorProtoToDenseTensorProto(s, model_path, tensor)); break; } default: @@ -732,6 +844,7 @@ struct GetElementSize { }; common::Status SparseTensorProtoToDenseTensorProto(const ONNX_NAMESPACE::SparseTensorProto& sparse, + const Path& model_path, ONNX_NAMESPACE::TensorProto& dense) { Status status = Status::OK(); @@ -758,7 +871,7 @@ common::Status SparseTensorProtoToDenseTensorProto(const ONNX_NAMESPACE::SparseT // need to read in sparse data first as it could be in a type specific field, in raw data, or in external data size_t sparse_bytes = 0; std::unique_ptr sparse_data_storage; - ORT_RETURN_IF_ERROR(UnpackInitializerData(sparse_values, sparse_data_storage, sparse_bytes)); + ORT_RETURN_IF_ERROR(UnpackInitializerData(sparse_values, model_path, sparse_data_storage, sparse_bytes)); void* sparse_data = sparse_data_storage.get(); size_t element_size = 0; // We want to this list to match the one used below in DenseTensorToSparseTensorProto() @@ -857,6 +970,7 @@ void CopyElement(void* dst, const void* src, int64_t dst_index, int64_t src_inde } common::Status DenseTensorToSparseTensorProto(const ONNX_NAMESPACE::TensorProto& dense_proto, + const Path& model_path, ONNX_NAMESPACE::SparseTensorProto& result) { ORT_ENFORCE(HasDataType(dense_proto), "Must have a valid data type"); @@ -881,7 +995,7 @@ common::Status DenseTensorToSparseTensorProto(const ONNX_NAMESPACE::TensorProto& size_t tensor_bytes_size = 0; std::unique_ptr dense_raw_data; - ORT_RETURN_IF_ERROR(UnpackInitializerData(dense_proto, dense_raw_data, tensor_bytes_size)); + ORT_RETURN_IF_ERROR(UnpackInitializerData(dense_proto, model_path, dense_raw_data, tensor_bytes_size)); size_t element_size = 0; MLTypeCallDispatcherRet type_disp(data_type); ORT_RETURN_IF_ERROR(type_disp.InvokeWithUnsupportedPolicy(element_size)); @@ -920,28 +1034,41 @@ template common::Status GetSizeInBytesFromTensorProto(const ONN size_t* out); template common::Status GetSizeInBytesFromTensorProto<0>(const ONNX_NAMESPACE::TensorProto& tensor_proto, size_t* out); -#define CASE_UNPACK(TYPE, ELEMENT_TYPE, DATA_SIZE) \ - case ONNX_NAMESPACE::TensorProto_DataType::TensorProto_DataType_##TYPE: { \ - size_t element_count = 0; \ - if (initializer.has_raw_data()) { \ - tensor_byte_size = initializer.raw_data().size(); \ - element_count = tensor_byte_size / sizeof(ELEMENT_TYPE); \ - } else { \ - element_count = initializer.DATA_SIZE(); \ - tensor_byte_size = element_count * sizeof(ELEMENT_TYPE); \ - } \ - unpacked_tensor.reset(new uint8_t[tensor_byte_size]); \ - return onnxruntime::utils::UnpackTensor( \ - initializer, \ - initializer.has_raw_data() ? initializer.raw_data().data() : nullptr, \ - initializer.has_raw_data() ? initializer.raw_data().size() : 0, \ - reinterpret_cast(unpacked_tensor.get()), element_count); \ - break; \ +#define CASE_UNPACK(TYPE, ELEMENT_TYPE, DATA_SIZE) \ + case ONNX_NAMESPACE::TensorProto_DataType::TensorProto_DataType_##TYPE: { \ + if (initializer.data_location() == TensorProto_DataLocation_EXTERNAL) { \ + ORT_RETURN_IF_ERROR(ReadExternalDataForTensor( \ + initializer, \ + model_path.IsEmpty() ? nullptr : model_path.ParentPath().ToPathString().c_str(), \ + unpacked_tensor, \ + tensor_byte_size)); \ + tensor_data_length = tensor_byte_size; \ + return Status::OK(); \ + } else { \ + size_t element_count = 0; \ + if (initializer.has_raw_data()) { \ + tensor_byte_size = initializer.raw_data().size(); \ + element_count = tensor_byte_size / sizeof(ELEMENT_TYPE); \ + } else { \ + element_count = initializer.DATA_SIZE(); \ + tensor_byte_size = element_count * sizeof(ELEMENT_TYPE); \ + } \ + tensor_data_length = tensor_byte_size; \ + unpacked_tensor.reset(new uint8_t[tensor_data_length]); \ + return onnxruntime::utils::UnpackTensor( \ + initializer, \ + initializer.has_raw_data() ? initializer.raw_data().data() : nullptr, \ + initializer.has_raw_data() ? initializer.raw_data().size() : 0, \ + reinterpret_cast(unpacked_tensor.get()), element_count); \ + } \ + break; \ } Status UnpackInitializerData(const onnx::TensorProto& initializer, + const Path& model_path, std::unique_ptr& unpacked_tensor, - size_t& tensor_byte_size) { + size_t& tensor_data_length) { + SafeInt tensor_byte_size = tensor_data_length; switch (initializer.data_type()) { CASE_UNPACK(FLOAT, float, float_data_size); CASE_UNPACK(DOUBLE, double, double_data_size); diff --git a/onnxruntime/core/framework/tensorprotoutils.h b/onnxruntime/core/framework/tensorprotoutils.h index 7efda74075..82028c1dad 100644 --- a/onnxruntime/core/framework/tensorprotoutils.h +++ b/onnxruntime/core/framework/tensorprotoutils.h @@ -7,7 +7,9 @@ #include #include "core/common/common.h" +#include "core/common/path.h" #include "core/common/status.h" +#include "core/framework/endian_utils.h" #include "core/framework/allocator.h" #include "core/framework/ml_value.h" #include "core/framework/mem_buffer.h" @@ -58,22 +60,33 @@ ONNXTensorElementDataType GetTensorElementType(const ONNX_NAMESPACE::TensorProto template common::Status GetSizeInBytesFromTensorProto(const ONNX_NAMESPACE::TensorProto& tensor_proto, size_t* out); -template -Status UnpackTensor(const ONNX_NAMESPACE::TensorProto& tensor, const void* raw_data, size_t raw_data_len, - /*out*/ T* p_data, size_t expected_size); - -// Convert the NodeProto from a Constant node into a TensorProto that can be used as an initializer +// Convert the AttributeProto from a Constant node into a TensorProto that can be used as an initializer +// If AttributeProto contains a TensorProto, this tensor proto is converted as is including the case when the +// the data location is external. i.e. it does not load the external data. +// However if AttributeProto contains SparseTensorProto then it converts the data into dense tensor proto +// (including loading external data when applicable). +// model_path is used for contructing full path for external_data common::Status ConstantNodeProtoToTensorProto(const ONNX_NAMESPACE::NodeProto& node, + const Path& model_path, ONNX_NAMESPACE::TensorProto& tensor); // Convert a SparseTensorProto to a dense TensorProto +// If the SparseTensorProto contains external data then it loads the data and converts to dense tensor proto +// The resulting TensorProto will contain the data as raw data. +// model_path is used for contructing full path for external_data common::Status SparseTensorProtoToDenseTensorProto(const ONNX_NAMESPACE::SparseTensorProto& sparse, + const Path& model_path, ONNX_NAMESPACE::TensorProto& dense); #if !defined(ORT_MINIMAL_BUILD) +// Convert a TensorProto to a SparseTensorProto +// If the tensorproto contains external data then it loads the data and converts to sparse tensor +// The resulting SparseTensorProto will contain the data as raw data +// model_path is used for contructing full path for external_data common::Status DenseTensorToSparseTensorProto(const ONNX_NAMESPACE::TensorProto& dense, + const Path& model_path, ONNX_NAMESPACE::SparseTensorProto& sparse); -#endif // !ORT_MINIMAL_BUILD +#endif // !ORT_MINIMAL_BUILD inline bool HasDimValue(const ONNX_NAMESPACE::TensorShapeProto_Dimension& dim) { return dim.value_case() == ONNX_NAMESPACE::TensorShapeProto_Dimension::kDimValue; @@ -109,6 +122,13 @@ inline bool HasRawData(const ONNX_NAMESPACE::TensorProto& ten_proto) { ten_proto.has_raw_data(); // XXX: Figure out how to do in proto3 } +inline bool HasExternalData(const ONNX_NAMESPACE::TensorProto& ten_proto) { + // Can not be UNDEFINED and can not be STRING but test for STRING is usually performed separately + // to return an error + return ten_proto.data_type() != ONNX_NAMESPACE::TensorProto::UNDEFINED && + ten_proto.data_location() == ONNX_NAMESPACE::TensorProto_DataLocation_EXTERNAL; +} + inline bool HasDataType(const ONNX_NAMESPACE::TensorProto& ten_proto) { return ten_proto.data_type() != ONNX_NAMESPACE::TensorProto::UNDEFINED; } @@ -126,10 +146,9 @@ inline bool HasElemType(const ONNX_NAMESPACE::TypeProto_SparseTensor& ten_proto) } inline bool HasName(const ONNX_NAMESPACE::SparseTensorProto& ten_proto) { - return ten_proto.values().has_name(); // XXX + return ten_proto.values().has_name(); // XXX } - inline bool HasKeyType(const ONNX_NAMESPACE::TypeProto_Map& map_proto) { return map_proto.key_type() != ONNX_NAMESPACE::TensorProto::UNDEFINED; } @@ -219,9 +238,37 @@ inline bool HasName(const ONNX_NAMESPACE::NodeProto& node_proto) { return node_proto.has_name(); } -// UnpackTensor from either raw data or the type specific data field. +#if !defined(ORT_MINIMAL_BUILD) +// Unpack tensor which contains external data. Uses the tensor_proto_dir to construct the full path for external data. +// If tensor_proto_dir == nullptr then uses the current directory instead. +// This function does not unpack string_data of a tensor template -Status UnpackTensor(const ONNX_NAMESPACE::TensorProto& tensor, /*out*/ T* p_data, size_t expected_size) { +Status UnpackTensorWithExternalData(const ONNX_NAMESPACE::TensorProto& tensor, + const ORTCHAR_T* tensor_proto_dir, size_t expected_size, + /*out*/ T* p_data); +#endif // !defined(ORT_MINIMAL_BUILD) + +// UnpackTensor from raw data or the type specific data field. Does not handle external data. +// If the tensor does not contain raw data then raw_data should be nullptr and raw_data_len should be 0. +template +Status UnpackTensor(const ONNX_NAMESPACE::TensorProto& tensor, const void* raw_data, size_t raw_data_len, + /*out*/ T* p_data, size_t expected_size); + +// UnpackTensor from raw data, external data or the type specific data field. +// Uses the model path to construct the full path for loading external data. In case when model_path is empty +// it uses current directory. +template +Status UnpackTensor(const ONNX_NAMESPACE::TensorProto& tensor, const Path& model_path, /*out*/ T* p_data, size_t expected_size) { +#if !defined(ORT_MINIMAL_BUILD) + if (HasExternalData(tensor)) { + auto tensor_proto_path = model_path.IsEmpty() ? nullptr : model_path.ParentPath().ToPathString().c_str(); + return UnpackTensorWithExternalData(tensor, tensor_proto_path, expected_size, p_data); + } +#else + ORT_UNUSED_PARAMETER(model_path); + ORT_RETURN_IF(HasExternalData(tensor), "TensorProto with external data is not supported in ORT minimal build."); +#endif + return HasRawData(tensor) ? UnpackTensor(tensor, tensor.raw_data().data(), tensor.raw_data().size(), p_data, expected_size) : UnpackTensor(tensor, nullptr, 0, p_data, expected_size); @@ -231,11 +278,13 @@ Status UnpackTensor(const ONNX_NAMESPACE::TensorProto& tensor, /*out*/ T* p_data * Unpack the data from an initializer tensor * Please note, this function does not unpack string_data of an initializer tensor * @param initializer given initializer tensor + * @param initializer_dir model_path to construct external data dir path. When this is empty, current dir is used. * @param unpacked_tensor the data from the initaizlier in uint8_t* form * @param tensor_byte_size the byte size of the unpacked_tensor * @returns Status::OK() if data is unpacked successfully */ common::Status UnpackInitializerData(const ONNX_NAMESPACE::TensorProto& initializer, + const Path& model_path, std::unique_ptr& unpacked_tensor, size_t& tensor_byte_size) ORT_MUST_USE_RESULT; diff --git a/onnxruntime/core/graph/graph.cc b/onnxruntime/core/graph/graph.cc index dd8708fa16..1f3de24a89 100644 --- a/onnxruntime/core/graph/graph.cc +++ b/onnxruntime/core/graph/graph.cc @@ -430,6 +430,10 @@ void Node::SetPriority(int priority) noexcept { priority_ = priority; } +const Path& Node::ModelPath() const noexcept { + return graph_->ModelPath(); +} + #if !defined(ORT_MINIMAL_BUILD) const Function* Node::GetFunctionBody(bool try_init_func_body) { @@ -966,6 +970,7 @@ Graph::Graph(const Model& owning_model, is_loaded_from_model_file_(GraphLoadedFromModelFile(graph_proto_)) { ORT_ENFORCE(graph_proto != nullptr, "graph_proto cannot be null"); ArgNameToTypeMap name_to_type_map; + const auto& model_path = ModelPath(); // Process 'Constant' nodes // Put the 'TensorProto' stored in the 'Constant' nodes attribute into the graphs initializer list @@ -975,7 +980,7 @@ Graph::Graph(const Model& owning_model, } const gsl::not_null tensor{graph_proto_->add_initializer()}; - auto status = utils::ConstantNodeProtoToTensorProto(node, *tensor); + auto status = utils::ConstantNodeProtoToTensorProto(node, model_path, *tensor); ORT_ENFORCE(status.IsOK(), status.ToString()); if (node.attribute(0).type() == AttributeProto_AttributeType_SPARSE_TENSOR) { auto p = sparse_tensor_names_.emplace(tensor->name()); @@ -1000,7 +1005,7 @@ Graph::Graph(const Model& owning_model, for (const auto& sparse_tensor : graph_proto_->sparse_initializer()) { ORT_ENFORCE(utils::HasName(sparse_tensor), "Sparse initializer must have a name. This model is invalid"); const gsl::not_null tensor{graph_proto_->add_initializer()}; - auto status = utils::SparseTensorProtoToDenseTensorProto(sparse_tensor, *tensor); + auto status = utils::SparseTensorProtoToDenseTensorProto(sparse_tensor, model_path, *tensor); ORT_ENFORCE(status.IsOK(), status.ToString()); auto p = sparse_tensor_names_.emplace(tensor->name()); ORT_ENFORCE(p.second, "Duplicate sparse_tensor_initializer: '", tensor->name(), "' Model is invalid."); @@ -2810,18 +2815,20 @@ common::Status Graph::SaveToOrtFormat(flatbuffers::FlatBufferBuilder& builder, std::vector> initializers_data; assert(sparse_tensor_names_.size() <= name_to_initial_tensor_.size()); initializers_data.reserve(name_to_initial_tensor_.size() - sparse_tensor_names_.size()); + const auto& model_path = ModelPath(); + for (const auto& pair : name_to_initial_tensor_) { if (sparse_tensor_names_.find(pair.first) == sparse_end) { flatbuffers::Offset fbs_tensor; ORT_RETURN_IF_ERROR( - experimental::utils::SaveInitializerOrtFormat(builder, *pair.second, fbs_tensor)); + experimental::utils::SaveInitializerOrtFormat(builder, *pair.second, model_path, fbs_tensor)); initializers_data.push_back(fbs_tensor); } else { SparseTensorProto sparse_initializer; - ORT_RETURN_IF_ERROR(utils::DenseTensorToSparseTensorProto(*pair.second, sparse_initializer)); + ORT_RETURN_IF_ERROR(utils::DenseTensorToSparseTensorProto(*pair.second, model_path, sparse_initializer)); flatbuffers::Offset fbs_sparse_tensor; ORT_RETURN_IF_ERROR( - experimental::utils::SaveSparseInitializerOrtFormat(builder, sparse_initializer, fbs_sparse_tensor)); + experimental::utils::SaveSparseInitializerOrtFormat(builder, sparse_initializer, model_path, fbs_sparse_tensor)); sparse_initializers_data.push_back(fbs_sparse_tensor); } } @@ -2995,6 +3002,10 @@ ONNX_NAMESPACE::GraphProto Graph::ToGraphProto() const { GraphProto result; ToGraphProtoInternal(result); + // Path of the owning model + // This is used for constructing full path for external data + // if it exists + const auto& model_path = ModelPath(); // We want to make sure that sparse initializers do not appear // as dense duplicates within the initializers list. @@ -3006,7 +3017,7 @@ ONNX_NAMESPACE::GraphProto Graph::ToGraphProto() const { *mutable_initializer->Add() = initializer; } else { auto& sparse_initializer = *result.add_sparse_initializer(); - auto status = utils::DenseTensorToSparseTensorProto(initializer, sparse_initializer); + auto status = utils::DenseTensorToSparseTensorProto(initializer, model_path, sparse_initializer); ORT_ENFORCE(status.IsOK(), "Failed to convert dense initializer to sparse"); } } @@ -3495,13 +3506,14 @@ Status Graph::InlineFunction(Node& node) { } RemoveNode(node.Index()); + const auto& model_path = ModelPath(); for (const auto& subgraph_node : subgraph.Nodes()) { if (subgraph_node.OpType() == kConstant) { // Copy constant nodes _value to name_to_initial_tensor_ ONNX_NAMESPACE::NodeProto subgraph_node_proto{}; subgraph_node.ToProto(subgraph_node_proto); const gsl::not_null tensor{graph_proto_->add_initializer()}; - ORT_RETURN_IF_ERROR(utils::ConstantNodeProtoToTensorProto(subgraph_node_proto, *tensor)); + ORT_RETURN_IF_ERROR(utils::ConstantNodeProtoToTensorProto(subgraph_node_proto, model_path, *tensor)); name_to_initial_tensor_[tensor->name()] = tensor; } else { std::vector inputs, outputs; @@ -3697,12 +3709,14 @@ common::Status Graph::LoadFromOrtFormat(const onnxruntime::experimental::fbs::Gr if (fbs_sparse_initializers) { sparse_tensor_names_.reserve(fbs_sparse_initializers->size()); + const auto& model_path = ModelPath(); + for (const auto* fbs_sparse_tensor : *fbs_sparse_initializers) { ORT_RETURN_IF(nullptr == fbs_sparse_tensor, "Sparse Initializer tensor is missing. Invalid ORT format model."); SparseTensorProto sparse_initializer; ORT_RETURN_IF_ERROR(experimental::utils::LoadSparseInitializerOrtFormat(*fbs_sparse_tensor, sparse_initializer)); TensorProto& initializer = *deserialized_proto_data_.add_initializer(); - ORT_RETURN_IF_ERROR(utils::SparseTensorProtoToDenseTensorProto(sparse_initializer, initializer)); + ORT_RETURN_IF_ERROR(utils::SparseTensorProtoToDenseTensorProto(sparse_initializer, model_path, initializer)); auto p = name_to_initial_tensor_.emplace(initializer.name(), &initializer); if (!p.second) { LOGS(logger_, WARNING) << "Duplicate initializer (dense, sparse or ConstantNode): '" << initializer.name() diff --git a/onnxruntime/core/graph/graph_flatbuffers_utils.cc b/onnxruntime/core/graph/graph_flatbuffers_utils.cc index 7f2087383d..7e59eca1bd 100644 --- a/onnxruntime/core/graph/graph_flatbuffers_utils.cc +++ b/onnxruntime/core/graph/graph_flatbuffers_utils.cc @@ -28,6 +28,7 @@ SaveDims(flatbuffers::FlatBufferBuilder& builder, const DimsFieldType& dims) { Status SaveInitializerOrtFormat(flatbuffers::FlatBufferBuilder& builder, const TensorProto& initializer, + const Path& model_path, flatbuffers::Offset& fbs_tensor) { auto name = SaveStringToOrtFormat(builder, initializer.has_name(), initializer.name()); auto doc_string = SaveStringToOrtFormat(builder, initializer.has_doc_string(), initializer.doc_string()); @@ -46,7 +47,7 @@ Status SaveInitializerOrtFormat(flatbuffers::FlatBufferBuilder& builder, std::unique_ptr unpacked_tensor; size_t tensor_byte_size = 0; ORT_RETURN_IF_ERROR( - onnxruntime::utils::UnpackInitializerData(initializer, unpacked_tensor, tensor_byte_size)); + onnxruntime::utils::UnpackInitializerData(initializer, model_path, unpacked_tensor, tensor_byte_size)); raw_data = builder.CreateVector(unpacked_tensor.get(), tensor_byte_size); } @@ -65,16 +66,17 @@ Status SaveInitializerOrtFormat(flatbuffers::FlatBufferBuilder& builder, Status SaveSparseInitializerOrtFormat(flatbuffers::FlatBufferBuilder& builder, const ONNX_NAMESPACE::SparseTensorProto& initializer, + const Path& model_path, flatbuffers::Offset& fbs_sparse_tensor) { // values const auto& values = initializer.values(); flatbuffers::Offset values_off; - ORT_RETURN_IF_ERROR(SaveInitializerOrtFormat(builder, values, values_off)); + ORT_RETURN_IF_ERROR(SaveInitializerOrtFormat(builder, values, model_path, values_off)); // Indicies const auto& indicies = initializer.indices(); flatbuffers::Offset indicies_off; - ORT_RETURN_IF_ERROR(SaveInitializerOrtFormat(builder, indicies, indicies_off)); + ORT_RETURN_IF_ERROR(SaveInitializerOrtFormat(builder, indicies, model_path, indicies_off)); // Shape auto shape = SaveDims(builder, initializer.dims()); @@ -122,7 +124,7 @@ Status SaveAttributeOrtFormat(flatbuffers::FlatBufferBuilder& builder, case fbs::AttributeType::TENSOR: { flatbuffers::Offset fbs_tensor; ORT_RETURN_IF_ERROR( - experimental::utils::SaveInitializerOrtFormat(builder, attr_proto.t(), fbs_tensor)); + experimental::utils::SaveInitializerOrtFormat(builder, attr_proto.t(), graph->ModelPath(), fbs_tensor)); GET_FBS_ATTR(builder, type, t, fbs_tensor); } break; case fbs::AttributeType::GRAPH: { @@ -152,7 +154,7 @@ Status SaveAttributeOrtFormat(flatbuffers::FlatBufferBuilder& builder, for (const auto& tensor : attr_proto.tensors()) { flatbuffers::Offset fbs_tensor; ORT_RETURN_IF_ERROR( - experimental::utils::SaveInitializerOrtFormat(builder, tensor, fbs_tensor)); + experimental::utils::SaveInitializerOrtFormat(builder, tensor, graph->ModelPath(), fbs_tensor)); fbs_tensors_vec.push_back(fbs_tensor); } auto tensors = builder.CreateVector(fbs_tensors_vec); diff --git a/onnxruntime/core/graph/graph_flatbuffers_utils.h b/onnxruntime/core/graph/graph_flatbuffers_utils.h index fe3ae9334b..e67430d330 100644 --- a/onnxruntime/core/graph/graph_flatbuffers_utils.h +++ b/onnxruntime/core/graph/graph_flatbuffers_utils.h @@ -19,9 +19,10 @@ namespace onnxruntime { class Graph; class Node; +class Path; namespace logging { -class Logger; + class Logger; } namespace experimental { @@ -36,11 +37,11 @@ namespace utils { // TODO, add ORT_MUST_USE_RESULT when it is moved to a different header onnxruntime::common::Status SaveInitializerOrtFormat( flatbuffers::FlatBufferBuilder& builder, const ONNX_NAMESPACE::TensorProto& initializer, - flatbuffers::Offset& fbs_tensor); + const Path& model_path, flatbuffers::Offset& fbs_tensor); onnxruntime::common::Status SaveSparseInitializerOrtFormat( flatbuffers::FlatBufferBuilder& builder, const ONNX_NAMESPACE::SparseTensorProto& initializer, - flatbuffers::Offset& fbs_sparse_tensor); + const Path& model_path, flatbuffers::Offset& fbs_sparse_tensor); // Convert a given AttributeProto into fbs::Attribute // Note, we current do not support graphs, and sparse_tensor(s) diff --git a/onnxruntime/core/optimizer/matmul_scale_fusion.cc b/onnxruntime/core/optimizer/matmul_scale_fusion.cc index 955393f964..bd577324d8 100644 --- a/onnxruntime/core/optimizer/matmul_scale_fusion.cc +++ b/onnxruntime/core/optimizer/matmul_scale_fusion.cc @@ -17,9 +17,9 @@ namespace onnxruntime { namespace { template struct ExtractScalarAsFloatDispatchTarget { - Status operator()(const ONNX_NAMESPACE::TensorProto& tensor_proto, float& scalar_float) { + Status operator()(const ONNX_NAMESPACE::TensorProto& tensor_proto, const Path& model_path, float& scalar_float) { T scalar; - ORT_RETURN_IF_ERROR(utils::UnpackTensor(tensor_proto, &scalar, 1)); + ORT_RETURN_IF_ERROR(utils::UnpackTensor(tensor_proto, model_path, &scalar, 1)); scalar_float = static_cast(scalar); return Status::OK(); } @@ -48,7 +48,7 @@ optional GetScalarConstantInitializer(const Graph& graph, const NodeArg& Status, ExtractScalarAsFloatDispatchTarget, uint32_t, uint64_t, int32_t, int64_t, MLFloat16, float, double, BFloat16> dispatcher{initializer->data_type()}; - ORT_THROW_IF_ERROR(dispatcher.Invoke(*initializer, scalar)); + ORT_THROW_IF_ERROR(dispatcher.Invoke(*initializer, graph.ModelPath(), scalar)); return {scalar}; } diff --git a/onnxruntime/core/providers/cpu/generator/constant_of_shape.cc b/onnxruntime/core/providers/cpu/generator/constant_of_shape.cc index f307cddd38..dd5596c581 100644 --- a/onnxruntime/core/providers/cpu/generator/constant_of_shape.cc +++ b/onnxruntime/core/providers/cpu/generator/constant_of_shape.cc @@ -41,6 +41,7 @@ void onnxruntime::ConstantOfShapeBase::SetValueFromTensorProto(const ONNX_NAMESP using namespace utils; ORT_ENFORCE(utils::HasDataType(t_proto)); ORT_ENFORCE(TensorProto::DataType_IsValid(t_proto.data_type())); + ORT_ENFORCE(!utils::HasExternalData(t_proto), "Tensor proto with external data for value attribute is not supported."); const auto tensor_type = static_cast(t_proto.data_type()); const void* const raw_data = utils::HasRawData(t_proto) ? t_proto.raw_data().data() : nullptr; const size_t raw_data_len = utils::HasRawData(t_proto) ? t_proto.raw_data().size() : 0; @@ -48,7 +49,7 @@ void onnxruntime::ConstantOfShapeBase::SetValueFromTensorProto(const ONNX_NAMESP case TensorProto::BOOL: FETCH_VALUE_DATA(bool); break; - case TensorProto::FLOAT: + case TensorProto::FLOAT: FETCH_VALUE_DATA(float); break; case TensorProto::FLOAT16: @@ -89,14 +90,13 @@ void onnxruntime::ConstantOfShapeBase::SetValueFromTensorProto(const ONNX_NAMESP #undef FETCH_VALUE_DATA - template inline void FilloutOutput(T value, void* output_data, size_t size) { auto out = gsl::make_span(reinterpret_cast(output_data), size); std::fill(out.begin(), out.end(), value); } -ConstantOfShapeBase::ConstantOfShapeBase(const OpKernelInfo& info){ +ConstantOfShapeBase::ConstantOfShapeBase(const OpKernelInfo& info) { TensorProto t_proto; if (info.GetAttr("value", &t_proto).IsOK()) { ORT_ENFORCE(t_proto.dims_size() == 1, "Must have a single dimension"); @@ -128,7 +128,6 @@ Status ConstantOfShapeBase::PrepareCompute(OpKernelContext* ctx, Tensor** output } Status ConstantOfShape::Compute(OpKernelContext* ctx) const { - Tensor* output_tensor = nullptr; ORT_RETURN_IF_ERROR(PrepareCompute(ctx, &output_tensor)); diff --git a/onnxruntime/core/providers/cpu/tensor/eye_like.cc b/onnxruntime/core/providers/cpu/tensor/eye_like.cc index 06bd30bcef..b0031c8291 100644 --- a/onnxruntime/core/providers/cpu/tensor/eye_like.cc +++ b/onnxruntime/core/providers/cpu/tensor/eye_like.cc @@ -2,7 +2,6 @@ // Licensed under the MIT License. #include "core/providers/cpu/tensor/eye_like.h" -#include "core/framework/tensorprotoutils.h" #include "core/util/math_cpuonly.h" using namespace ::onnxruntime::common; diff --git a/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/helper.cc b/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/helper.cc index 55739d12df..9ec99afb96 100644 --- a/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/helper.cc +++ b/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/helper.cc @@ -231,7 +231,8 @@ bool HasValidQuantizationZeroPoints(const InitializedTensorSet& initializers, co std::unique_ptr unpacked_tensor; size_t tensor_byte_size; - auto status = onnxruntime::utils::UnpackInitializerData(zero_tensor, unpacked_tensor, tensor_byte_size); + auto status = onnxruntime::utils::UnpackInitializerData(zero_tensor, node.ModelPath(), + unpacked_tensor, tensor_byte_size); if (!status.IsOK()) { LOGS_DEFAULT(ERROR) << "QLinearConv erro when unpack zero tensor:" << status.ErrorMessage(); return false; diff --git a/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/op_builder.cc b/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/op_builder.cc index 1e5badf13f..58d8c9b2a4 100644 --- a/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/op_builder.cc +++ b/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/op_builder.cc @@ -291,7 +291,8 @@ static Status AddInitializerInNewLayout(ModelBuilder& model_builder, case ONNX_NAMESPACE::TensorProto_DataType_UINT8: case ONNX_NAMESPACE::TensorProto_DataType_INT8: { ORT_RETURN_IF_ERROR( - onnxruntime::utils::UnpackInitializerData(tensor, unpacked_tensor, tensor_byte_size)); + onnxruntime::utils::UnpackInitializerData(tensor, model_builder.GetGraphViewer().ModelPath(), + unpacked_tensor, tensor_byte_size)); src = unpacked_tensor.get(); break; } @@ -371,7 +372,8 @@ static Status AddInitializerTransposed(ModelBuilder& model_builder, case ONNX_NAMESPACE::TensorProto_DataType_UINT8: case ONNX_NAMESPACE::TensorProto_DataType_INT8: { ORT_RETURN_IF_ERROR( - onnxruntime::utils::UnpackInitializerData(tensor, unpacked_tensor, tensor_byte_size)); + onnxruntime::utils::UnpackInitializerData(tensor, model_builder.GetGraphViewer().ModelPath(), + unpacked_tensor, tensor_byte_size)); src = unpacked_tensor.get(); break; } @@ -498,7 +500,8 @@ static Status GetQuantizationZeroPoint(const ModelBuilder& model_builder, const size_t tensor_byte_size; const auto& zero_point_tensor = *model_builder.GetInitializerTensors().at(node.InputDefs()[idx]->Name()); ORT_RETURN_IF_ERROR( - onnxruntime::utils::UnpackInitializerData(zero_point_tensor, unpacked_tensor, tensor_byte_size)); + onnxruntime::utils::UnpackInitializerData(zero_point_tensor, model_builder.GetGraphViewer().ModelPath(), + unpacked_tensor, tensor_byte_size)); zero_point = static_cast(unpacked_tensor.get()[0]); return Status::OK(); } diff --git a/onnxruntime/test/framework/sparse_kernels_test.cc b/onnxruntime/test/framework/sparse_kernels_test.cc index 3e72ba1a5d..332f585f76 100644 --- a/onnxruntime/test/framework/sparse_kernels_test.cc +++ b/onnxruntime/test/framework/sparse_kernels_test.cc @@ -14,6 +14,7 @@ #include "core/session/inference_session.h" #include "test/providers/provider_test_utils.h" #include "test_utils.h" +#include "file_util.h" #include "gtest/gtest.h" #include "gmock/gmock.h" @@ -517,6 +518,27 @@ std::vector CreateValues() { } */ +template +static void CreateTensorWithExternalData( + TensorProto_DataType type, + const std::vector& test_data, + std::basic_string& filename, + TensorProto& tensor_proto) { + // Create external data + FILE* fp; + CreateTestFile(fp, filename); + size_t size_in_bytes = test_data.size() * sizeof(T); + ASSERT_EQ(size_in_bytes, fwrite(test_data.data(), 1, size_in_bytes, fp)); + ASSERT_EQ(0, fclose(fp)); + + // set the tensor_proto to reference this external data + onnx::StringStringEntryProto* location = tensor_proto.mutable_external_data()->Add(); + location->set_key("location"); + location->set_value(ToMBString(filename)); + tensor_proto.set_data_location(onnx::TensorProto_DataLocation_EXTERNAL); + tensor_proto.set_data_type(type); +} + template static NodeProto CreateConstantNode(bool indices_1D, std::function& values, TensorProto& tp)> inserter, @@ -577,7 +599,9 @@ static void TestConversion(bool use_1D_indices, auto node = CreateConstantNode(use_1D_indices, inserter, expected); TensorProto dense; - utils::ConstantNodeProtoToTensorProto(node, dense); + // Path is required for loading external data (if any) + // When path is empty it will look for the data in current dir + utils::ConstantNodeProtoToTensorProto(node, Path(), dense); gsl::span expected_span = gsl::make_span(expected.data(), expected.size()); checker(expected_span, dense); @@ -646,6 +670,15 @@ TEST(SparseTensorConversionTests, TestConstantNodeConversion) { RawDataWriter(values, tp, TensorProto_DataType_UINT8); }, RawDataChecker); + + // Test constant node conversion for SparseTensor with external data + PathString tensor_filename(ORT_TSTR("tensor_XXXXXX")); + TestConversion(true, + [&tensor_filename](const std::vector& values, TensorProto& tp) { + CreateTensorWithExternalData(TensorProto_DataType_FLOAT, values, tensor_filename, tp); + }, + RawDataChecker); + DeleteFileFromDisk(tensor_filename.c_str()); } /// Dense to Sparse conversion tests @@ -736,10 +769,13 @@ static void TestDenseToSparseConversion( checker) { std::vector expected_values; std::vector expected_indicies; + // Path is required for loading external data + // Using empty path here since the data is not external + Path model_path; TensorProto dense_tensor = CreateDenseTensor(inserter, expected_values, expected_indicies); SparseTensorProto sparse_tensor; - utils::DenseTensorToSparseTensorProto(dense_tensor, sparse_tensor); + utils::DenseTensorToSparseTensorProto(dense_tensor, model_path, sparse_tensor); gsl::span expected_values_span = gsl::make_span(expected_values.data(), expected_values.size()); diff --git a/onnxruntime/test/framework/tensorutils_test.cc b/onnxruntime/test/framework/tensorutils_test.cc index 633b6159a1..a048df52fc 100644 --- a/onnxruntime/test/framework/tensorutils_test.cc +++ b/onnxruntime/test/framework/tensorutils_test.cc @@ -4,10 +4,15 @@ #include "core/framework/tensorprotoutils.h" #include "core/graph/onnx_protobuf.h" #include "test/util/include/asserts.h" +#include "file_util.h" #include "gtest/gtest.h" #include "gmock/gmock.h" +#ifdef _WIN32 +#include +#endif + using namespace ::onnxruntime::utils; using namespace ONNX_NAMESPACE; @@ -16,7 +21,7 @@ namespace test { //T must be float for double, and it must match with the 'type' argument template -void test_unpack_float_tensor(TensorProto_DataType type) { +void TestUnpackFloatTensor(TensorProto_DataType type, const Path& model_path) { TensorProto float_tensor_proto; float_tensor_proto.set_data_type(type); T f[4] = {1.1f, 2.2f, 3.3f, 4.4f}; @@ -27,7 +32,7 @@ void test_unpack_float_tensor(TensorProto_DataType type) { } float_tensor_proto.set_raw_data(rawdata, len); T float_data2[4]; - auto status = UnpackTensor(float_tensor_proto, float_data2, 4); + auto status = UnpackTensor(float_tensor_proto, model_path, float_data2, 4); EXPECT_TRUE(status.IsOK()) << status.ErrorMessage(); EXPECT_EQ(1.1f, float_data2[0]); EXPECT_EQ(2.2f, float_data2[1]); @@ -35,22 +40,26 @@ void test_unpack_float_tensor(TensorProto_DataType type) { EXPECT_EQ(4.4f, float_data2[3]); } -TEST(TensorParseTest, TensorUtilsTest) { +TEST(TensorProtoUtilsTest, UnpackTensor) { TensorProto bool_tensor_proto; + // Path is required for loading external data. + // Using empty path here since this test does not test + // external data utils + Path model_path; bool_tensor_proto.set_data_type(TensorProto_DataType_BOOL); bool_tensor_proto.add_int32_data(1); bool bool_data[1]; - auto status = UnpackTensor(bool_tensor_proto, bool_data, 1); + auto status = UnpackTensor(bool_tensor_proto, model_path, bool_data, 1); EXPECT_TRUE(status.IsOK()) << status.ErrorMessage(); EXPECT_TRUE(bool_data[0]); float float_data[1]; - status = UnpackTensor(bool_tensor_proto, float_data, 1); + status = UnpackTensor(bool_tensor_proto, model_path, float_data, 1); EXPECT_FALSE(status.IsOK()); - test_unpack_float_tensor(TensorProto_DataType_FLOAT); - test_unpack_float_tensor(TensorProto_DataType_DOUBLE); + TestUnpackFloatTensor(TensorProto_DataType_FLOAT, model_path); + TestUnpackFloatTensor(TensorProto_DataType_DOUBLE, model_path); TensorProto string_tensor_proto; string_tensor_proto.set_data_type(TensorProto_DataType_STRING); @@ -58,12 +67,12 @@ TEST(TensorParseTest, TensorUtilsTest) { string_tensor_proto.add_string_data("b"); std::string string_data[2]; - status = UnpackTensor(string_tensor_proto, string_data, 2); + status = UnpackTensor(string_tensor_proto, model_path, string_data, 2); EXPECT_TRUE(status.IsOK()) << status.ErrorMessage(); EXPECT_EQ("a", string_data[0]); EXPECT_EQ("b", string_data[1]); - status = UnpackTensor(bool_tensor_proto, string_data, 2); + status = UnpackTensor(bool_tensor_proto, model_path, string_data, 2); EXPECT_FALSE(status.IsOK()); } @@ -77,6 +86,56 @@ std::vector CreateValues() { return {"one", "two", "three", "four"}; } +template +static void CreateTensorWithExternalData( + TensorProto_DataType type, + const std::vector& test_data, + std::basic_string& filename, + TensorProto& tensor_proto) { + // Create external data + FILE* fp; + CreateTestFile(fp, filename); + size_t size_in_bytes = test_data.size() * sizeof(T); + ASSERT_EQ(size_in_bytes, fwrite(test_data.data(), 1, size_in_bytes, fp)); + ASSERT_EQ(0, fclose(fp)); + + // set the tensor_proto to reference this external data + onnx::StringStringEntryProto* location = tensor_proto.mutable_external_data()->Add(); + location->set_key("location"); + location->set_value(ToMBString(filename)); + tensor_proto.mutable_dims()->Add(test_data.size()); + tensor_proto.set_data_location(onnx::TensorProto_DataLocation_EXTERNAL); + tensor_proto.set_data_type(type); +} + +template +static void TestUnpackExternalTensor(TensorProto_DataType type, const Path& model_path) { + // Create external data + std::basic_string filename(ORT_TSTR("tensor_XXXXXX")); + TensorProto tensor_proto; + auto test_data = CreateValues(); + CreateTensorWithExternalData(type, test_data, filename, tensor_proto); + std::unique_ptr file_deleter(const_cast(filename.c_str()), + DeleteFileFromDisk); + + // Unpack tensor with external data + std::vector val(test_data.size()); + auto st = utils::UnpackTensor(tensor_proto, model_path, val.data(), test_data.size()); + ASSERT_TRUE(st.IsOK()) << st.ErrorMessage(); + + // Validate data + for (size_t i = 0; i < test_data.size(); i++) { + ASSERT_EQ(val[i], test_data[i]); + } +} + +TEST(TensorProtoUtilsTest, UnpackTensorWithExternalData) { + Path model_path; + TestUnpackExternalTensor(TensorProto_DataType_FLOAT, model_path); + TestUnpackExternalTensor(TensorProto_DataType_DOUBLE, model_path); + TestUnpackExternalTensor(TensorProto_DataType_INT32, model_path); +} + template static NodeProto CreateConstantNode(const std::string& attrib_name, AttributeProto_AttributeType type, std::function add_data) { @@ -111,7 +170,8 @@ static void TestConstantNodeConversion(const std::string& attrib_name, [&input, &add_data](AttributeProto& attrib) { add_data(attrib, input); }); TensorProto tp; - EXPECT_STATUS_OK(utils::ConstantNodeProtoToTensorProto(c, tp)); + Path model_path; + EXPECT_STATUS_OK(utils::ConstantNodeProtoToTensorProto(c, model_path, tp)); EXPECT_THAT(get_data(tp), ::testing::ContainerEq(input)); } @@ -174,5 +234,50 @@ TEST(TensorProtoUtilsTest, ConstantTensorProto) { // sparse_tensor is covered by SparseTensorConversionTests.TestConstantNodeConversion } + +template +static NodeProto CreateConstantNodeWithExternalData(TensorProto_DataType type, PathString& tensor_filename, + const std::vector& test_data) { + NodeProto constant_node; + constant_node.set_op_type("Constant"); + constant_node.add_output("Constant_output"); + + AttributeProto& attrib = *constant_node.mutable_attribute()->Add(); + attrib.set_name("attrib"); + attrib.set_type(AttributeProto_AttributeType_TENSOR); + + TensorProto& tp = *attrib.mutable_t(); + CreateTensorWithExternalData(type, test_data, tensor_filename, tp); + + return constant_node; +} + +template +static void TestConstantNodeConversionWithExternalData(TensorProto_DataType type) { + // Create a constant node with external data + auto test_data = CreateValues(); + Path model_path; + PathString tensor_filename(ORT_TSTR("tensor_XXXXXX")); + auto c = CreateConstantNodeWithExternalData(type, tensor_filename, test_data); + std::unique_ptr file_deleter(const_cast(tensor_filename.c_str()), + DeleteFileFromDisk); + + // Convert NodeProto to tensorproto (with external data) + TensorProto tp; + EXPECT_STATUS_OK(utils::ConstantNodeProtoToTensorProto(c, model_path, tp)); + + // Unpack tensor and validate the data + std::vector val(test_data.size()); + auto st = utils::UnpackTensor(tp, model_path, val.data(), test_data.size()); + ASSERT_TRUE(st.IsOK()) << st.ErrorMessage(); + for (size_t i = 0; i < test_data.size(); i++) { + ASSERT_EQ(val[i], test_data[i]); + } +} + +TEST(TensorProtoUtilsTest, ConstantTensorProtoWithExternalData) { + TestConstantNodeConversionWithExternalData(TensorProto_DataType_FLOAT); + TestConstantNodeConversionWithExternalData(TensorProto_DataType_DOUBLE); +} } // namespace test } // namespace onnxruntime