/* Copyright 2015 The TensorFlow Authors. All Rights Reserved.
|
|
Licensed under the Apache License, Version 2.0 (the "License");
|
you may not use this file except in compliance with the License.
|
You may obtain a copy of the License at
|
|
http://www.apache.org/licenses/LICENSE-2.0
|
|
Unless required by applicable law or agreed to in writing, software
|
distributed under the License is distributed on an "AS IS" BASIS,
|
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
See the License for the specific language governing permissions and
|
limitations under the License.
|
==============================================================================*/
|
|
#include "tensorflow/core/util/tensor_slice_set.h"
|
|
#include <vector>
|
#include "tensorflow/core/lib/core/errors.h"
|
#include "tensorflow/core/lib/gtl/map_util.h"
|
#include "tensorflow/core/platform/logging.h"
|
#include "tensorflow/core/util/tensor_slice_util.h"
|
|
namespace tensorflow {
|
|
namespace checkpoint {
|
|
TensorSliceSet::TensorSliceSet(const TensorShape& shape, DataType type)
|
: shape_(shape), type_(type) {}
|
|
TensorSliceSet::~TensorSliceSet() {}
|
|
Status TensorSliceSet::Register(const TensorSlice& slice, const string& tag,
|
const float* data) {
|
TensorShape result_shape;
|
TF_RETURN_IF_ERROR(slice.SliceTensorShape(shape_, &result_shape));
|
string str = slice.DebugString();
|
|
if (slices_.empty()) {
|
slices_hull_ = slice;
|
} else {
|
// We check if there is any intersection between this slice and any of the
|
// registered slices.
|
if (slices_hull_.Overlaps(slice)) {
|
for (const auto& x : slices_) {
|
if (slice.Overlaps(x.second.slice)) {
|
return errors::Internal("Overlapping slices: existing slice = ",
|
x.first, ", new slice = ", str);
|
}
|
}
|
}
|
// No overlap: we can now insert the slice
|
slices_hull_.UpdateToCover(slice);
|
}
|
|
TensorSliceSet::SliceInfo info = {slice, tag, data,
|
result_shape.num_elements()};
|
slices_.insert(std::make_pair(str, info));
|
return Status::OK();
|
}
|
|
// TODO(yangke): merge Query() with QueryMeta()
|
bool TensorSliceSet::Query(const TensorSlice& slice, float* data) const {
|
Status s;
|
string str = slice.DebugString();
|
// First we check if there is an exactly match (this is the dominant case).
|
const TensorSliceSet::SliceInfo* info = gtl::FindOrNull(slices_, str);
|
if (info) {
|
if (data) {
|
std::copy_n(info->data, info->num_floats, data);
|
}
|
return true;
|
} else {
|
// We didn't find any exact match but there is still a possibility that
|
// multiple existing slices can be patched together to output the slice.
|
// We figure this out by computing the intersection of each of the existing
|
// slices with the query slice, and check if the union of all these
|
// intersections cover the entire slice. We rely on the fact that the
|
// existing slices don't have any intersection among themselves.
|
TensorShape target_shape;
|
Status s;
|
s = slice.SliceTensorShape(shape_, &target_shape);
|
if (!s.ok()) {
|
LOG(WARNING) << s;
|
return false;
|
}
|
int64 total_size = target_shape.num_elements();
|
|
int64 overlap_size = 0;
|
TensorSlice intersection;
|
TensorShape inter_shape;
|
for (const auto& x : slices_) {
|
if (slice.Intersect(x.second.slice, &intersection)) {
|
s = intersection.SliceTensorShape(shape_, &inter_shape);
|
if (!s.ok()) {
|
LOG(WARNING) << s;
|
return false;
|
}
|
overlap_size += inter_shape.num_elements();
|
}
|
}
|
if (total_size == overlap_size) {
|
// We have it!
|
// Now we need to copy the data to "data"
|
if (data) {
|
for (const auto& x : slices_) {
|
CopyDataFromTensorSliceToTensorSlice(shape_, x.second.slice, slice,
|
x.second.data, data);
|
}
|
}
|
return true;
|
} else {
|
// We don't have all the data for the asked tensor slice
|
return false;
|
}
|
}
|
}
|
|
bool TensorSliceSet::QueryMeta(
|
const TensorSlice& slice,
|
std::vector<std::pair<TensorSlice, string>>* results) const {
|
results->clear();
|
Status s;
|
string str = slice.DebugString();
|
// First we check if there is an exactly match (this is the dominant case).
|
const TensorSliceSet::SliceInfo* info = gtl::FindOrNull(slices_, str);
|
if (info) {
|
results->emplace_back(std::make_pair(info->slice, info->tag));
|
return true;
|
} else {
|
// We didn't find any exact match but there is still a possibility that
|
// multiple existing slices can be patched together to output the slice.
|
// We figure this out by computing the intersection of each of the existing
|
// slices with the query slice, and check if the union of all these
|
// intersections cover the entire slice. We rely on the fact that the
|
// existing slices don't have any intersection among themselves.
|
TensorShape target_shape;
|
Status s;
|
s = slice.SliceTensorShape(shape_, &target_shape);
|
if (!s.ok()) {
|
LOG(WARNING) << s;
|
return false;
|
}
|
int64 total_size = target_shape.num_elements();
|
|
int64 overlap_size = 0;
|
TensorSlice intersection;
|
TensorShape inter_shape;
|
for (const auto& x : slices_) {
|
if (slice.Intersect(x.second.slice, &intersection)) {
|
s = intersection.SliceTensorShape(shape_, &inter_shape);
|
if (!s.ok()) {
|
LOG(WARNING) << s;
|
return false;
|
}
|
overlap_size += inter_shape.num_elements();
|
results->emplace_back(std::make_pair(x.second.slice, x.second.tag));
|
}
|
}
|
if (total_size == overlap_size) {
|
// We have it!
|
return true;
|
} else {
|
// We don't have all the data for the asked tensor slice
|
results->clear();
|
return false;
|
}
|
}
|
}
|
|
Status RegisterTensorSlice(
|
const string& name, const TensorShape& shape, DataType type,
|
const string& tag, const TensorSlice& slice,
|
std::unordered_map<string, TensorSliceSet*>* tensor_slices) {
|
DCHECK_NE(tensor_slices, nullptr);
|
TensorSliceSet* tss = gtl::FindPtrOrNull(*tensor_slices, name);
|
// Create a tensor slice set if needed
|
if (!tss) {
|
tss = new TensorSliceSet(shape, type);
|
tensor_slices->insert(std::make_pair(name, tss));
|
} else {
|
// Check if the shapes match
|
const TensorShape& tss_shape(tss->shape());
|
if (!shape.IsSameSize(tss_shape)) {
|
return errors::Internal("Incompatible tensor shapes detected for tensor ",
|
name, ": existing = ", tss_shape.DebugString(),
|
", new = ", shape.DebugString());
|
}
|
if (type != tss->type()) {
|
return errors::Internal("Incompatible tensor types detected for tensor ",
|
name,
|
": existing = ", DataTypeString(tss->type()),
|
", new = ", DataTypeString(type));
|
}
|
}
|
// Register the tensor slices without the actual data.
|
return tss->Register(slice, tag, nullptr);
|
}
|
|
} // namespace checkpoint
|
|
} // namespace tensorflow
|
