1 /* Copyright 2015 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 // The utility to write checkpoints for google brain tensor ops and v3 17 // checkpoints for dist_belief. 18 19 #ifndef TENSORFLOW_UTIL_TENSOR_SLICE_WRITER_H_ 20 #define TENSORFLOW_UTIL_TENSOR_SLICE_WRITER_H_ 21 22 #include <unordered_map> 23 24 #include "tensorflow/core/framework/tensor_shape.h" 25 #include "tensorflow/core/framework/tensor_slice.h" 26 #include "tensorflow/core/framework/types.h" 27 #include "tensorflow/core/lib/core/errors.h" 28 #include "tensorflow/core/lib/core/status.h" 29 #include "tensorflow/core/lib/core/stringpiece.h" 30 #include "tensorflow/core/lib/gtl/map_util.h" 31 #include "tensorflow/core/lib/strings/stringprintf.h" 32 #include "tensorflow/core/platform/logging.h" 33 #include "tensorflow/core/platform/macros.h" 34 #include "tensorflow/core/platform/types.h" 35 #include "tensorflow/core/util/saved_tensor_slice.pb_text.h" 36 #include "tensorflow/core/util/saved_tensor_slice.pb.h" 37 #include "tensorflow/core/util/saved_tensor_slice_util.h" 38 39 namespace tensorflow { 40 41 namespace checkpoint { 42 43 class TensorSliceWriter { 44 public: 45 // Abstract interface that TensorSliceWriter uses for building 46 class Builder { 47 public: 48 virtual ~Builder() {} 49 virtual void Add(StringPiece key, StringPiece value) = 0; 50 virtual Status Finish(int64* file_size) = 0; 51 }; 52 typedef std::function<Status(const string&, Builder**)> CreateBuilderFunction; 53 54 TensorSliceWriter(const string& filename, 55 CreateBuilderFunction create_builder); 56 virtual ~TensorSliceWriter() {} 57 // Adds a slice. We support float and int32 for now. 58 // TODO(yangke): add more supports 59 template <typename T> 60 Status Add(const string& name, const TensorShape& shape, 61 const TensorSlice& slice, const T* data); 62 Status Finish(); 63 64 // Allocate "num_elements" elements in "ss" and save the data in "data" 65 // there. 66 template <typename T> 67 static Status SaveData(const T* data, int64 num_elements, SavedSlice* ss); 68 69 static size_t MaxBytesPerElement(DataType dt); 70 71 private: 72 static const size_t kMaxMessageBytes = 1LL << 31; 73 // Filling in the TensorProto in a SavedSlice will add the following 74 // header bytes, in addition to the data: 75 // - 1 byte: TensorProto tag and wire format 76 // - <= 5 bytes: TensorProto length 77 // - 1 byte: Repeated *_val tag and wire format 78 // - <= 5 bytes: *_val length 79 // However, we add 1KB of slack, to be conservative and guard 80 // against other additions to the TensorProto. 81 static const size_t kTensorProtoHeaderBytes = 1 << 10; 82 83 const string filename_; 84 const CreateBuilderFunction create_builder_; 85 const string tmpname_; 86 87 // A mapping from the tensor names to their index in meta_.saved_slice_meta() 88 std::unordered_map<string, int> name_to_index_; 89 // The metadata that holds all the saved tensor slices. 90 SavedTensorSlices sts_; 91 // The data to be written to the builder 92 std::map<string, string> data_; 93 // Total number of slices written 94 int slices_; 95 TF_DISALLOW_COPY_AND_ASSIGN(TensorSliceWriter); 96 }; 97 98 template <typename T> 99 Status TensorSliceWriter::Add(const string& name, const TensorShape& shape, 100 const TensorSlice& slice, const T* data) { 101 // The tensor and the slice have to be compatible 102 if (shape.dims() != slice.dims()) { 103 return errors::Internal("Incompatible tensor shape and slice: ", "shape = ", 104 shape.DebugString(), 105 ", slice = ", slice.DebugString()); 106 } 107 DataType dt = DataTypeToEnum<T>::value; 108 // We need to add an entry for "name" if there isn't an entry already. 109 int index = gtl::FindWithDefault(name_to_index_, name, -1); 110 if (index >= 0) { 111 // The same tensor has been registered -- we verify that the shapes and the 112 // type agree. 113 const SavedSliceMeta& ssm = sts_.meta().tensor(index); 114 CHECK_EQ(name, ssm.name()) << ProtoShortDebugString(ssm); 115 TensorShape ssm_shape(ssm.shape()); 116 if (!shape.IsSameSize(ssm_shape)) { 117 return errors::Internal( 118 "Mismatching shapes: existing tensor = ", ssm_shape.DebugString(), 119 ", trying to add name ", name, ", shape = ", shape.DebugString()); 120 } 121 if (dt != ssm.type()) { 122 return errors::Internal( 123 "Mismatching types: existing type = ", DataTypeString(ssm.type()), 124 ", trying to add name ", name, ", type = ", DataTypeString(dt)); 125 } 126 } else { 127 // Insert the new tensor name with the shape information 128 index = sts_.meta().tensor_size(); 129 name_to_index_.insert(std::make_pair(name, index)); 130 SavedSliceMeta* ssm = sts_.mutable_meta()->add_tensor(); 131 ssm->set_name(name); 132 shape.AsProto(ssm->mutable_shape()); 133 ssm->set_type(dt); 134 } 135 // Now we need to add the slice info the list of slices. 136 SavedSliceMeta* ssm = sts_.mutable_meta()->mutable_tensor(index); 137 slice.AsProto(ssm->add_slice()); 138 139 // Now we need to add the real data. 140 { 141 SavedTensorSlices sts; 142 SavedSlice* ss = sts.mutable_data(); 143 ss->set_name(name); 144 slice.AsProto(ss->mutable_slice()); 145 TensorShape saved_shape(ssm->shape()); 146 TensorShape sliced_shape; 147 TF_RETURN_IF_ERROR(slice.SliceTensorShape(saved_shape, &sliced_shape)); 148 TF_RETURN_IF_ERROR(SaveData(data, sliced_shape.num_elements(), ss)); 149 string key = EncodeTensorNameSlice(name, slice); 150 // TODO(yangke): consider doing a two-pass thing where the first pass just 151 // list the tensor slices we want to save and then another pass to actually 152 // set the data. Need to figure out if the interface works well. 153 std::pair<string, string> key_value(key, ""); 154 if (!sts.AppendToString(&key_value.second)) { 155 return errors::Internal("Error writing Tensor. Possible size overflow."); 156 } 157 data_.insert(key_value); 158 } 159 ++slices_; 160 return Status::OK(); 161 } 162 163 template <typename T> 164 Status TensorSliceWriter::SaveData(const T* data, int64 num_elements, 165 SavedSlice* ss) { 166 size_t size_bound = 167 ss->ByteSize() + kTensorProtoHeaderBytes + 168 (MaxBytesPerElement(DataTypeToEnum<T>::value) * num_elements); 169 if (size_bound > kMaxMessageBytes) { 170 return errors::InvalidArgument( 171 "Tensor slice is too large to serialize (conservative estimate: ", 172 size_bound, " bytes)"); 173 } 174 Fill(data, num_elements, ss->mutable_data()); 175 DCHECK_GE(ss->ByteSize(), 0); 176 DCHECK_LE(ss->ByteSize(), size_bound); 177 return Status::OK(); 178 } 179 180 template <> 181 Status TensorSliceWriter::SaveData(const string* data, int64 num_elements, 182 SavedSlice* ss); 183 184 // Create a table builder that will write to "filename" in 185 // tensorflow::io::Table format. If successful, return OK 186 // and set "*builder" to the allocated builder. Otherwise, return a 187 // non-OK status. 188 Status CreateTableTensorSliceBuilder(const string& filename, 189 TensorSliceWriter::Builder** builder); 190 191 } // namespace checkpoint 192 193 } // namespace tensorflow 194 195 #endif // TENSORFLOW_UTIL_TENSOR_SLICE_WRITER_H_ 196
