1 /* Copyright 2017 The TensorFlow Authors. All Rights Reserved. 2 3 Licensed under the Apache License, Version 2.0 (the "License"); 4 you may not use this file except in compliance with the License. 5 You may obtain a copy of the License at 6 7 http://www.apache.org/licenses/LICENSE-2.0 8 9 Unless required by applicable law or agreed to in writing, software 10 distributed under the License is distributed on an "AS IS" BASIS, 11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 See the License for the specific language governing permissions and 13 limitations under the License. 14 ==============================================================================*/ 15 16 #ifndef TENSORFLOW_CORE_FRAMEWORK_VARIANT_TENSOR_DATA_H_ 17 #define TENSORFLOW_CORE_FRAMEWORK_VARIANT_TENSOR_DATA_H_ 18 19 #include <algorithm> 20 #include <vector> 21 22 #include "tensorflow/core/framework/tensor.h" 23 #include "tensorflow/core/lib/core/stringpiece.h" 24 #include "tensorflow/core/platform/types.h" 25 26 namespace tensorflow { 27 28 class VariantTensorDataProto; 29 30 // The serialization format for Variant objects. Objects with references to 31 // other Tensors can simply store those tensors in the `tensors` field, and 32 // serialize other metadata content in to the `metadata` field. Objects can 33 // optionally set the `type_name` for type-checking before deserializing an 34 // object. 35 // 36 // This is the native C++ class equivalent of VariantTensorDataProto. They are 37 // separate so that kernels do not need to depend on protos. 38 class VariantTensorData { 39 public: 40 VariantTensorData() = default; 41 42 // TODO(b/118823936): This silently returns if the proto is invalid. 43 // Consider calling FromProto explicitly instead. 44 VariantTensorData(VariantTensorDataProto proto); 45 46 // Name of the type of objects being serialized. type_name()47 const string& type_name() const { return type_name_; } set_type_name(const string & type_name)48 void set_type_name(const string& type_name) { type_name_ = type_name; } 49 50 template <typename T, bool = std::is_pod<typename std::decay<T>::type>::value> 51 struct PODResolver {}; 52 53 // Portions of the object that are not Tensors. 54 // Directly supported types include string POD types. 55 template <typename T> set_metadata(const T & value)56 void set_metadata(const T& value) { 57 SetMetadata<T>(value, PODResolver<T>()); 58 } 59 60 template <typename T> get_metadata(T * value)61 bool get_metadata(T* value) const { 62 return GetMetadata<T>(value, PODResolver<T>()); 63 } 64 65 // Tensors contained within objects being serialized. 66 int tensors_size() const; 67 const Tensor& tensors(int index) const; 68 const std::vector<Tensor>& tensors() const; 69 Tensor* add_tensors(); 70 71 // A more general version of add_tensors. Parameters are perfectly forwarded 72 // to the constructor of the tensor added here. 73 template <typename... TensorConstructorArgs> 74 Tensor* add_tensor(TensorConstructorArgs&&... args); 75 76 // Conversion to and from VariantTensorDataProto 77 void ToProto(VariantTensorDataProto* proto) const; 78 // This allows optimizations via std::move. 79 bool FromProto(VariantTensorDataProto proto); 80 bool FromConstProto(const VariantTensorDataProto& proto); 81 82 // Serialization via VariantTensorDataProto 83 string SerializeAsString() const; 84 bool SerializeToString(string* buf); 85 bool ParseFromString(string s); 86 87 string DebugString() const; 88 89 public: 90 string type_name_; 91 string metadata_; 92 std::vector<Tensor> tensors_; 93 94 private: 95 template <typename T> SetMetadata(const string & value,PODResolver<T,false>)96 void SetMetadata(const string& value, PODResolver<T, false /* is_pod */>) { 97 metadata_ = value; 98 } 99 100 template <typename T> GetMetadata(string * value,PODResolver<T,false>)101 bool GetMetadata(string* value, PODResolver<T, false /* is_pod */>) const { 102 *value = metadata_; 103 return true; 104 } 105 106 template <typename T> SetMetadata(const T & value,PODResolver<T,true>)107 void SetMetadata(const T& value, PODResolver<T, true /* is_pod */>) { 108 metadata_.assign(reinterpret_cast<const char*>(&value), sizeof(T)); 109 } 110 111 template <typename T> GetMetadata(T * value,PODResolver<T,true>)112 bool GetMetadata(T* value, PODResolver<T, true /* is_pod */>) const { 113 if (metadata_.size() != sizeof(T)) return false; 114 std::copy_n(metadata_.data(), sizeof(T), reinterpret_cast<char*>(value)); 115 return true; 116 } 117 }; 118 119 // For backwards compatibility for when this was a proto 120 string ProtoDebugString(const VariantTensorData& object); 121 122 } // namespace tensorflow 123 124 #endif // TENSORFLOW_CORE_FRAMEWORK_VARIANT_TENSOR_DATA_H_ 125