From a555740708906c54dca4053d1446e2e4da1a002c Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Thu, 28 Oct 2021 10:38:10 -0700 Subject: [PATCH] Attention fusion: update uint8 tensor parsing for ONNX upgrade (#9564) * use UnpackTensor to parse uint8 tensor * address review feedback --- .../core/optimizer/attention_fusion_helper.h | 20 +++++++++++++++++-- 1 file changed, 18 insertions(+), 2 deletions(-) diff --git a/onnxruntime/core/optimizer/attention_fusion_helper.h b/onnxruntime/core/optimizer/attention_fusion_helper.h index 57b912d624..df7bb4e11a 100644 --- a/onnxruntime/core/optimizer/attention_fusion_helper.h +++ b/onnxruntime/core/optimizer/attention_fusion_helper.h @@ -2,6 +2,7 @@ // Licensed under the MIT License. #include "onnx/defs/shape_inference.h" #include "onnx/defs/tensor_proto_util.h" +#include "core/framework/tensorprotoutils.h" #pragma once @@ -322,8 +323,23 @@ bool ValidateUnidirMask(const Graph& graph, const NodeArg& mask, bool& is_unidir } if (tensor_proto->data_type() == ONNX_NAMESPACE::TensorProto_DataType_UINT8) { - std::vector int32_data = ONNX_NAMESPACE::ParseData(tensor_proto); - if (!ValidateUnidirMask(int32_data, shape->dim(2).dim_value(), is_unidirectional)) { + size_t bytes; + if (!utils::GetSizeInBytesFromTensorProto<0>(*tensor_proto, &bytes).IsOK()) { + return false; + } + auto data = std::make_unique(bytes); + uint8_t* p = data.get(); + if (!utils::UnpackTensor( + *tensor_proto, + tensor_proto->raw_data().size() ? tensor_proto->raw_data().data() : nullptr, + tensor_proto->raw_data().size(), + p, + bytes) + .IsOK()) { + return false; + } + std::vector mask_data(p, p + bytes); + if (!ValidateUnidirMask(mask_data, shape->dim(2).dim_value(), is_unidirectional)) { DEBUG_LOG("Mask is neither unidirectional nor all ones"); return false; }