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_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:
~Builder()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);
~TensorSliceWriter()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>
Add(const string & name,const TensorShape & shape,const TensorSlice & slice,const T * data)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>
SaveData(const T * data,int64 num_elements,SavedSlice * ss)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_CORE_UTIL_TENSOR_SLICE_WRITER_H_
196