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