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