Rework TensorToTensorProto. Do not put string data to raw_string. Eliminate redundant argument. (#3438)

Rework TensorToTensorProto. Eliminate redundant argument.
  Do not put string data into raw_data.
This commit is contained in:
Dmitri Smirnov 2020-04-07 11:42:10 -07:00 committed by GitHub
parent 43d6c464fc
commit 53b9d52fc6
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
3 changed files with 15 additions and 14 deletions

View file

@ -518,8 +518,7 @@ ONNXTensorElementDataType GetTensorElementType(const ONNX_NAMESPACE::TensorProto
return CApiElementTypeFromProtoType(tensor_proto.data_type());
}
ONNX_NAMESPACE::TensorProto TensorToTensorProto(const Tensor& tensor, const std::string& tensor_proto_name,
const ONNX_NAMESPACE::TypeProto& tensor_proto_type) {
ONNX_NAMESPACE::TensorProto TensorToTensorProto(const Tensor& tensor, const std::string& tensor_proto_name) {
// Given we are using the raw_data field in the protobuf, this will work only for little-endian format.
ORT_ENFORCE(endian::native == endian::little);
@ -532,12 +531,17 @@ ONNX_NAMESPACE::TensorProto TensorToTensorProto(const Tensor& tensor, const std:
tensor_proto.add_dims(dim);
}
// TODO Once utils::GetTensorProtoType supports all data types, you can get the tensor proto type from the tensor,
// as follows (which will allow us to get rid of the tensor_proto_type argument).
//tensor_proto.set_data_type(utils::GetTensorProtoType(tensor));
tensor_proto.set_data_type(tensor_proto_type.tensor_type().elem_type());
tensor_proto.set_raw_data(tensor.DataRaw(), tensor.SizeInBytes());
tensor_proto.set_data_type(tensor.GetElementType());
if (tensor.IsDataTypeString()) {
auto* mutable_string_data = tensor_proto.mutable_string_data();
auto f = tensor.Data<std::string>();
auto end = f + tensor.Shape().Size();
for (; f < end; ++f) {
*mutable_string_data->Add() = *f;
}
} else {
tensor_proto.set_raw_data(tensor.DataRaw(), tensor.SizeInBytes());
}
return tensor_proto;
}

View file

@ -38,14 +38,11 @@ common::Status TensorProtoToMLValue(const Env& env, const ORTCHAR_T* tensor_prot
/** Creates a TensorProto from a Tensor.
@param[in] tensor the Tensor whose data and shape will be used to create the TensorProto.
@param[in] tensor_proto_name the name of the TensorProto.
@param[in] tensor_proto_type the type of the TensorProto.
@return the TensorProto.
Note: Method currently requires that data is in little-endian format.
TODO Once the GetTensorProtoType supports all data types, we can remove the tensor_proto_type parameter and
instead get the type from the tensor. */
ONNX_NAMESPACE::TensorProto TensorToTensorProto(const Tensor& tensor, const std::string& tensor_proto_name,
const ONNX_NAMESPACE::TypeProto& tensor_proto_type);
*/
ONNX_NAMESPACE::TensorProto TensorToTensorProto(const Tensor& tensor, const std::string& tensor_proto_name);
ONNXTensorElementDataType CApiElementTypeFromProtoType(int type);
ONNXTensorElementDataType GetTensorElementType(const ONNX_NAMESPACE::TensorProto& tensor_proto);

View file

@ -94,7 +94,7 @@ Status ConstantFolding::ApplyImpl(Graph& graph, bool& modified, int graph_level,
ORT_ENFORCE(ort_value.IsTensor());
const Tensor& out_tensor = ort_value.Get<Tensor>();
ONNX_NAMESPACE::TensorProto out_tensorproto =
utils::TensorToTensorProto(out_tensor, constant_arg_out->Name(), *constant_arg_out->TypeAsProto());
utils::TensorToTensorProto(out_tensor, constant_arg_out->Name());
graph.AddInitializedTensor(out_tensorproto);
}