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 #include "tensorflow/core/util/tensor_slice_set.h"
17
18 #include <vector>
19 #include "tensorflow/core/lib/core/errors.h"
20 #include "tensorflow/core/lib/gtl/map_util.h"
21 #include "tensorflow/core/platform/logging.h"
22 #include "tensorflow/core/util/tensor_slice_util.h"
23
24 namespace tensorflow {
25
26 namespace checkpoint {
27
TensorSliceSet(const TensorShape & shape,DataType type)28 TensorSliceSet::TensorSliceSet(const TensorShape& shape, DataType type)
29 : shape_(shape), type_(type) {}
30
~TensorSliceSet()31 TensorSliceSet::~TensorSliceSet() {}
32
Register(const TensorSlice & slice,const string & tag,const float * data)33 Status TensorSliceSet::Register(const TensorSlice& slice, const string& tag,
34 const float* data) {
35 TensorShape result_shape;
36 TF_RETURN_IF_ERROR(slice.SliceTensorShape(shape_, &result_shape));
37 string str = slice.DebugString();
38
39 if (slices_.empty()) {
40 slices_hull_ = slice;
41 } else {
42 // We check if there is any intersection between this slice and any of the
43 // registered slices.
44 if (slices_hull_.Overlaps(slice)) {
45 for (const auto& x : slices_) {
46 if (slice.Overlaps(x.second.slice)) {
47 return errors::Internal("Overlapping slices: existing slice = ",
48 x.first, ", new slice = ", str);
49 }
50 }
51 }
52 // No overlap: we can now insert the slice
53 slices_hull_.UpdateToCover(slice);
54 }
55
56 TensorSliceSet::SliceInfo info = {slice, tag, data,
57 result_shape.num_elements()};
58 slices_.insert(std::make_pair(str, info));
59 return Status::OK();
60 }
61
62 // TODO(yangke): merge Query() with QueryMeta()
Query(const TensorSlice & slice,float * data) const63 bool TensorSliceSet::Query(const TensorSlice& slice, float* data) const {
64 Status s;
65 string str = slice.DebugString();
66 // First we check if there is an exactly match (this is the dominant case).
67 const TensorSliceSet::SliceInfo* info = gtl::FindOrNull(slices_, str);
68 if (info) {
69 if (data) {
70 std::copy_n(info->data, info->num_floats, data);
71 }
72 return true;
73 } else {
74 // We didn't find any exact match but there is still a possibility that
75 // multiple existing slices can be patched together to output the slice.
76 // We figure this out by computing the intersection of each of the existing
77 // slices with the query slice, and check if the union of all these
78 // intersections cover the entire slice. We rely on the fact that the
79 // existing slices don't have any intersection among themselves.
80 TensorShape target_shape;
81 Status s;
82 s = slice.SliceTensorShape(shape_, &target_shape);
83 if (!s.ok()) {
84 LOG(WARNING) << s;
85 return false;
86 }
87 int64 total_size = target_shape.num_elements();
88
89 int64 overlap_size = 0;
90 TensorSlice intersection;
91 TensorShape inter_shape;
92 for (const auto& x : slices_) {
93 if (slice.Intersect(x.second.slice, &intersection)) {
94 s = intersection.SliceTensorShape(shape_, &inter_shape);
95 if (!s.ok()) {
96 LOG(WARNING) << s;
97 return false;
98 }
99 overlap_size += inter_shape.num_elements();
100 }
101 }
102 if (total_size == overlap_size) {
103 // We have it!
104 // Now we need to copy the data to "data"
105 if (data) {
106 for (const auto& x : slices_) {
107 CopyDataFromTensorSliceToTensorSlice(shape_, x.second.slice, slice,
108 x.second.data, data);
109 }
110 }
111 return true;
112 } else {
113 // We don't have all the data for the asked tensor slice
114 return false;
115 }
116 }
117 }
118
QueryMeta(const TensorSlice & slice,std::vector<std::pair<TensorSlice,string>> * results) const119 bool TensorSliceSet::QueryMeta(
120 const TensorSlice& slice,
121 std::vector<std::pair<TensorSlice, string>>* results) const {
122 results->clear();
123 Status s;
124 string str = slice.DebugString();
125 // First we check if there is an exactly match (this is the dominant case).
126 const TensorSliceSet::SliceInfo* info = gtl::FindOrNull(slices_, str);
127 if (info) {
128 results->emplace_back(std::make_pair(info->slice, info->tag));
129 return true;
130 } else {
131 // We didn't find any exact match but there is still a possibility that
132 // multiple existing slices can be patched together to output the slice.
133 // We figure this out by computing the intersection of each of the existing
134 // slices with the query slice, and check if the union of all these
135 // intersections cover the entire slice. We rely on the fact that the
136 // existing slices don't have any intersection among themselves.
137 TensorShape target_shape;
138 Status s;
139 s = slice.SliceTensorShape(shape_, &target_shape);
140 if (!s.ok()) {
141 LOG(WARNING) << s;
142 return false;
143 }
144 int64 total_size = target_shape.num_elements();
145
146 int64 overlap_size = 0;
147 TensorSlice intersection;
148 TensorShape inter_shape;
149 for (const auto& x : slices_) {
150 if (slice.Intersect(x.second.slice, &intersection)) {
151 s = intersection.SliceTensorShape(shape_, &inter_shape);
152 if (!s.ok()) {
153 LOG(WARNING) << s;
154 return false;
155 }
156 overlap_size += inter_shape.num_elements();
157 results->emplace_back(std::make_pair(x.second.slice, x.second.tag));
158 }
159 }
160 if (total_size == overlap_size) {
161 // We have it!
162 return true;
163 } else {
164 // We don't have all the data for the asked tensor slice
165 results->clear();
166 return false;
167 }
168 }
169 }
170
RegisterTensorSlice(const string & name,const TensorShape & shape,DataType type,const string & tag,const TensorSlice & slice,std::unordered_map<string,TensorSliceSet * > * tensor_slices)171 Status RegisterTensorSlice(
172 const string& name, const TensorShape& shape, DataType type,
173 const string& tag, const TensorSlice& slice,
174 std::unordered_map<string, TensorSliceSet*>* tensor_slices) {
175 DCHECK_NE(tensor_slices, nullptr);
176 TensorSliceSet* tss = gtl::FindPtrOrNull(*tensor_slices, name);
177 // Create a tensor slice set if needed
178 if (!tss) {
179 tss = new TensorSliceSet(shape, type);
180 tensor_slices->insert(std::make_pair(name, tss));
181 } else {
182 // Check if the shapes match
183 const TensorShape& tss_shape(tss->shape());
184 if (!shape.IsSameSize(tss_shape)) {
185 return errors::Internal("Incompatible tensor shapes detected for tensor ",
186 name, ": existing = ", tss_shape.DebugString(),
187 ", new = ", shape.DebugString());
188 }
189 if (type != tss->type()) {
190 return errors::Internal("Incompatible tensor types detected for tensor ",
191 name,
192 ": existing = ", DataTypeString(tss->type()),
193 ", new = ", DataTypeString(type));
194 }
195 }
196 // Register the tensor slices without the actual data.
197 return tss->Register(slice, tag, nullptr);
198 }
199
200 } // namespace checkpoint
201
202 } // namespace tensorflow
203