include/onnx/proto_utils.h

0001 // Copyright (c) ONNX Project Contributors
0002 //
0003 // SPDX-License-Identifier: Apache-2.0
0004
0005 #pragma once
0006
0007 #include <google/protobuf/io/coded_stream.h>
0008 #include <google/protobuf/io/zero_copy_stream_impl_lite.h>
0009
0010 #include <string>
0011 #include <vector>
0012
0013 #include "onnx/onnx_pb.h"
0014
0015 #ifdef ONNX_USE_LITE_PROTO
0016 #include <google/protobuf/message_lite.h>
0017 #else // ONNX_USE_LITE_PROTO
0018 #include <google/protobuf/message.h>
0019 #endif // !ONNX_USE_LITE_PROTO
0020
0021 namespace ONNX_NAMESPACE {
0022
0023 #ifdef ONNX_USE_LITE_PROTO
0024 using ::google::protobuf::MessageLite;
0025 inline std::string ProtoDebugString(const MessageLite& proto) {
0026   // Since the MessageLite interface does not support reflection, there is very
0027   // little information that this and similar methods can provide.
0028   // But when using lite proto this is the best we can provide.
0029   return proto.ShortDebugString();
0030 }
0031 #else
0032 using ::google::protobuf::Message;
0033 inline std::string ProtoDebugString(const Message& proto) {
0034   return proto.ShortDebugString();
0035 }
0036 #endif
0037
0038 template <typename Proto>
0039 bool ParseProtoFromBytes(Proto* proto, const char* buffer, size_t length) {
0040   ::google::protobuf::io::ArrayInputStream input_stream(buffer, static_cast<int>(length));
0041   ::google::protobuf::io::CodedInputStream coded_stream(&input_stream);
0042   int total_bytes_limit = (2048LL << 20) - 1;
0043 #if GOOGLE_PROTOBUF_VERSION >= 3011000
0044   // Only take one parameter since protobuf 3.11
0045   coded_stream.SetTotalBytesLimit(total_bytes_limit);
0046 #else
0047   // Total bytes hard limit / warning limit are set to 2GB and 512MB respectively.
0048   coded_stream.SetTotalBytesLimit(total_bytes_limit, 512LL << 20);
0049 #endif
0050
0051   return proto->ParseFromCodedStream(&coded_stream);
0052 }
0053
0054 template <typename T>
0055 inline std::vector<T> RetrieveValues(const AttributeProto& attr);
0056 template <>
0057 inline std::vector<int64_t> RetrieveValues(const AttributeProto& attr) {
0058   return {attr.ints().begin(), attr.ints().end()};
0059 }
0060
0061 template <>
0062 inline std::vector<std::string> RetrieveValues(const AttributeProto& attr) {
0063   return {attr.strings().begin(), attr.strings().end()};
0064 }
0065
0066 template <>
0067 inline std::vector<float> RetrieveValues(const AttributeProto& attr) {
0068   return {attr.floats().begin(), attr.floats().end()};
0069 }
0070
0071 } // namespace ONNX_NAMESPACE