| /* 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/framework/op_kernel.h" |
| #include "tensorflow/core/framework/register_types.h" |
| #include "tensorflow/core/framework/tensor.h" |
| #include "tensorflow/core/framework/types.h" |
| #include "tensorflow/core/framework/variant.h" |
| #include "tensorflow/core/framework/variant_encode_decode.h" |
| #include "tensorflow/core/lib/gtl/inlined_vector.h" |
| |
| namespace tensorflow { |
| |
| namespace { |
| |
| class DeserializeSparseOp : public OpKernel { |
| public: |
| explicit DeserializeSparseOp(OpKernelConstruction* context) |
| : OpKernel(context) { |
| OP_REQUIRES_OK(context, context->GetAttr("dtype", &dtype_)); |
| } |
| |
| void Compute(OpKernelContext* context) override { |
| const Tensor& input = context->input(0); |
| |
| OP_REQUIRES( |
| context, input.dims() > 0, |
| errors::InvalidArgument("Serialized sparse should have non-zero rank ", |
| input.shape().DebugString())); |
| OP_REQUIRES(context, input.shape().dim_size(input.dims() - 1) == 3, |
| errors::InvalidArgument( |
| "Serialized sparse should have 3 as the last dimension ", |
| input.shape().DebugString())); |
| |
| // `input_dims_to_stack` is the number of dimensions that will be added to |
| // each of the elements before they are concatenated into the output. |
| const int64 input_dims_to_stack = input.dims() - 1; |
| int num_sparse_tensors = 1; |
| for (int i = 0; i < input_dims_to_stack; ++i) { |
| num_sparse_tensors *= input.shape().dim_size(i); |
| } |
| |
| if (num_sparse_tensors == 1 && input_dims_to_stack == 0) { |
| // Special case with a single sparse tensor, and no dimensions to add |
| // to the output indices. We can return the boxed tensors directly (after |
| // validating them). |
| const Tensor* output_indices; |
| const Tensor* output_values; |
| const Tensor* output_shape; |
| const auto& input_as_vec = input.vec<Variant>(); |
| int64 total_non_zeros; |
| OP_REQUIRES_OK(context, GetAndValidateSparseTensorShape( |
| input_as_vec(1), input_as_vec(2), 0, |
| &output_shape, &total_non_zeros)); |
| OP_REQUIRES_OK(context, GetAndValidateSparseTensorIndicesAndValues( |
| input_as_vec(0), input_as_vec(1), 0, |
| output_shape->NumElements(), &output_indices, |
| &output_values)); |
| context->set_output(0, *output_indices); |
| context->set_output(1, *output_values); |
| context->set_output(2, *output_shape); |
| return; |
| } |
| |
| OP_REQUIRES( |
| context, num_sparse_tensors > 0, |
| errors::InvalidArgument( |
| "Serialized sparse should have at least 1 serialized tensor, " |
| "but has a zero dimension ", |
| input.shape().DebugString())); |
| |
| const auto& input_as_matrix = input.flat_inner_dims<Variant, 2>(); |
| |
| // Compute the output "dense shape" of and number of non-zero elements in |
| // the stacked sparse tensors. Given an input of shape (S_0, ..., |
| // S_{input_dims_to_stack-1}, 3), and an element of dense shape (E_0, ... |
| // E_n), the output dense shape will be (S_0, ..., |
| // S_{input_dims_to_stack-1}, E_0, ..., E_n). |
| Tensor* output_shape; |
| int64 total_non_zeros = 0; |
| |
| // Allocate and build the initial output shape based on the element shape of |
| // the 0th sparse tensor in the input. |
| // |
| // NOTE(mrry): We define `element_shape` as a `const Tensor*` rather than a |
| // `Tensor` to avoid the overhead of allocating and deallocating a `Tensor` |
| // on the stack. While the per-`Tensor` cost is small, this op can unbox a |
| // large number of tensors (3 per batch element) and these fixed overheads |
| // dominate when the number of non-zeros per element is small. |
| const Tensor* element_shape; |
| OP_REQUIRES_OK(context, GetAndValidateSparseTensorShape( |
| input_as_matrix(0, 1), input_as_matrix(0, 2), 0, |
| &element_shape, &total_non_zeros)); |
| OP_REQUIRES_OK(context, |
| context->allocate_output( |
| 2, {input_dims_to_stack + element_shape->NumElements()}, |
| &output_shape)); |
| const auto element_shape_vec = element_shape->vec<int64>(); |
| auto output_shape_vec = output_shape->vec<int64>(); |
| output_shape_vec(0) = num_sparse_tensors; |
| for (int64 j = 0; j < input_dims_to_stack; ++j) { |
| output_shape_vec(j) = input.dim_size(j); |
| } |
| for (int64 j = 0; j < element_shape->NumElements(); ++j) { |
| output_shape_vec(j + input_dims_to_stack) = element_shape_vec(j); |
| } |
| |
| // Accumulate the number of non-zero elements from the remaining sparse |
| // tensors, and validate that they have compatible dense shapes. |
| // |
| // NOTE(mrry): For compatibility with the implementations of |
| // DeserializeManySparse, and many ops that generate SparseTensors to batch |
| // that do not have a fixed dense_shape (e.g. `tf.parse_single_example()`), |
| // we compute the maximum in each dimension to find the smallest dense_shape |
| // that bounds all of the input SparseTensors. |
| for (int i = 1; i < num_sparse_tensors; ++i) { |
| int64 num_non_zeros; |
| OP_REQUIRES_OK(context, GetAndValidateSparseTensorShape( |
| input_as_matrix(i, 1), input_as_matrix(i, 2), |
| i, &element_shape, &num_non_zeros)); |
| total_non_zeros += num_non_zeros; |
| OP_REQUIRES( |
| context, |
| output_shape->NumElements() - input_dims_to_stack == |
| element_shape->NumElements(), |
| errors::InvalidArgument( |
| "Inconsistent shape across SparseTensors: rank prior to " |
| "SparseTensor[", |
| i, "] was: ", output_shape->NumElements() - input_dims_to_stack, |
| " but rank of SparseTensor[", i, |
| "] is: ", element_shape->NumElements())); |
| const auto element_shape_vec = element_shape->vec<int64>(); |
| for (int j = 0; j < element_shape->NumElements(); ++j) { |
| output_shape_vec(j + input_dims_to_stack) = std::max( |
| output_shape_vec(j + input_dims_to_stack), element_shape_vec(j)); |
| } |
| } |
| |
| // Compute the output "indices" matrix and "values" vector. |
| Tensor* output_indices; |
| Tensor* output_values; |
| |
| const int output_rank = output_shape->NumElements(); |
| OP_REQUIRES_OK(context, |
| context->allocate_output( |
| 0, {static_cast<int64>(total_non_zeros), output_rank}, |
| &output_indices)); |
| OP_REQUIRES_OK( |
| context, context->allocate_output( |
| 1, {static_cast<int64>(total_non_zeros)}, &output_values)); |
| |
| // The bulk of the work in this method involves building the output indices |
| // in a tight loop. For cache friendliness, we generate the indices in the |
| // order that they will be laid out in memory. We use raw pointers instead |
| // of Eigen element/slice indexing methods, to access the underlying index |
| // buffer to minimize the amount of work in that tight loop. |
| int64* output_indices_data = output_indices->matrix<int64>().data(); |
| size_t current_row = 0; |
| |
| for (int i = 0; i < num_sparse_tensors; ++i) { |
| const Tensor* element_indices; |
| const Tensor* element_values; |
| OP_REQUIRES_OK(context, this->GetAndValidateSparseTensorIndicesAndValues( |
| input_as_matrix(i, 0), input_as_matrix(i, 1), |
| i, output_rank - input_dims_to_stack, |
| &element_indices, &element_values)); |
| |
| const size_t num_index_rows = element_values->NumElements(); |
| |
| // An empty sparse tensor in the input will generate no data |
| // in the output. We short-circuit the rest of the iteration to avoid |
| // triggering assertions in the Eigen when manipulating empty tensors (or |
| // slices of tensors). |
| if (num_index_rows == 0) continue; |
| |
| const size_t start_row = current_row; |
| const size_t next_start_row = current_row + num_index_rows; |
| |
| // NOTE(mrry): If the element is a scalar SparseTensor, |
| // `element_indices` will be an empty tensor, and this pointer will not |
| // be valid. However, we will not dereference the pointer in that case, |
| // because `input_dims_to_stack == output_rank`. |
| const int64* element_indices_data = |
| element_indices->matrix<int64>().data(); |
| |
| // Build the submatrix of `output_indices` for the i^th sparse tensor |
| // in the input. |
| // |
| // Each row of `output_indices` comprises `input_dims_to_stack` indices |
| // based on the position of the i^th sparse tensor in the input tensor, |
| // followed by the indices from the corresponding row in |
| // `element_indices`. |
| if (input_dims_to_stack == 1 && output_rank == 2) { |
| // We specialize this case because the compiler can generate |
| // more efficient code when the number of indices for each element is |
| // known statically. Since the most common use of this op is to |
| // serialize batches of SparseTensors, and the most common source of |
| // SparseTensors is the `tf.parse_single_example()` op, which generates |
| // 1-D SparseTensors, we statically unroll the loop for the rank 2 |
| // output case. |
| for (; current_row < next_start_row; ++current_row) { |
| *output_indices_data++ = i; |
| *output_indices_data++ = *element_indices_data++; |
| } |
| } else { |
| // `sparse_tensor_index` is the tuple of indices that correspond to |
| // mapping the flat element index (`i`) back onto the stacked |
| // coordinates implied by the position of the i^th sparse tensor in the |
| // input tensor. |
| // |
| // We build `sparse_tensor_index` in reverse (innermost/minor dimension |
| // to outermost/major dimension). The `cumulative_product` represents |
| // the size of the inner subtensor for which `sparse_tensor_index` has |
| // already been built. |
| gtl::InlinedVector<int64, 4> sparse_tensor_index(input_dims_to_stack); |
| int cumulative_product = 1; |
| for (size_t j = 0; j < sparse_tensor_index.size(); ++j) { |
| size_t reverse_index = sparse_tensor_index.size() - j - 1; |
| sparse_tensor_index[reverse_index] = |
| (i / cumulative_product) % input.dim_size(reverse_index); |
| cumulative_product *= input.dim_size(reverse_index); |
| } |
| for (; current_row < next_start_row; ++current_row) { |
| for (int64 sparse_tensor_index_component : sparse_tensor_index) { |
| *output_indices_data++ = sparse_tensor_index_component; |
| } |
| for (size_t k = input_dims_to_stack; k < output_rank; ++k) { |
| *output_indices_data++ = *element_indices_data++; |
| } |
| } |
| } |
| |
| // Build the subvector of `output_values` for the i^th sparse tensor |
| // in the input. |
| // |
| // NOTE(mrry): There is a potential optimization here where we use a T* |
| // to represent the current position in `output_values`, but it would |
| // require some rejigging of the template parameters. |
| // NOTE(mrry): Another potential optimization: if we know that this |
| // operation consumes its input, we could std::move non-primitive elements |
| // into the output and avoid a copy. |
| Eigen::DSizes<Eigen::DenseIndex, 1> values_start(start_row); |
| Eigen::DSizes<Eigen::DenseIndex, 1> values_sizes(num_index_rows); |
| |
| #define HANDLE_TYPE(T) \ |
| case DataTypeToEnum<T>::value: { \ |
| output_values->vec<T>().slice(values_start, values_sizes) = \ |
| element_values->vec<T>(); \ |
| break; \ |
| } |
| switch (dtype_) { |
| TF_CALL_ALL_TYPES(HANDLE_TYPE); |
| TF_CALL_QUANTIZED_TYPES(HANDLE_TYPE); |
| #undef HANDLE_TYPE |
| default: |
| OP_REQUIRES_OK( |
| context, errors::Unimplemented( |
| "DeserializeSparse Unhandled data type: ", dtype_)); |
| } |
| } |
| } |
| |
| private: |
| Status GetAndValidateSparseTensorShape(const Variant& serialized_values, |
| const Variant& serialized_shape, |
| int index, const Tensor** output_shape, |
| int64* output_num_non_zeros) { |
| // Deserialize and validate the shape. |
| *output_shape = serialized_shape.get<Tensor>(); |
| if (*output_shape == nullptr) { |
| return errors::InvalidArgument( |
| "Could not get a tensor from serialized_sparse[", index, ", 2]"); |
| } |
| if ((*output_shape)->dtype() != DT_INT64) { |
| return errors::InvalidArgument( |
| "Expected serialized_sparse[", index, |
| ", 2] to be a vector of DT_INT64 but received dtype ", |
| DataTypeString((*output_shape)->dtype())); |
| } |
| if (!TensorShapeUtils::IsVector((*output_shape)->shape())) { |
| return errors::InvalidArgument( |
| "Expected serialized_sparse[", index, |
| ", 2] to be a shape vector but its shape is ", |
| (*output_shape)->shape().DebugString()); |
| } |
| *output_num_non_zeros = serialized_values.get<Tensor>()->NumElements(); |
| return Status::OK(); |
| } |
| |
| Status GetAndValidateSparseTensorIndicesAndValues( |
| const Variant& serialized_indices, const Variant& serialized_values, |
| int index, int expected_rank, const Tensor** output_indices, |
| const Tensor** output_values) { |
| // Deserialize and validate the indices. |
| *output_indices = serialized_indices.get<Tensor>(); |
| if (*output_indices == nullptr) { |
| return errors::InvalidArgument( |
| "Could not get a tensor from serialized_sparse[", index, ", 0]"); |
| } |
| if ((*output_indices)->dtype() != DT_INT64) { |
| return errors::InvalidArgument( |
| "Expected serialized_sparse[", index, |
| ", 0] to be a matrix of DT_INT64 but received dtype ", |
| DataTypeString((*output_indices)->dtype())); |
| } |
| if (!TensorShapeUtils::IsMatrix((*output_indices)->shape())) { |
| return errors::InvalidArgument( |
| "Expected serialized_sparse[", index, |
| ", 0] to represent an index matrix but received shape ", |
| (*output_indices)->shape().DebugString()); |
| } |
| int64 num_entries = (*output_indices)->dim_size(0); |
| int rank = (*output_indices)->dim_size(1); |
| if (rank != expected_rank) { |
| return errors::InvalidArgument( |
| "Expected column counts of SparseTensor[", index, |
| "].indices to match size of SparseTensor[", index, |
| "].shape but they do not: ", rank, " vs. ", expected_rank); |
| } |
| |
| // Deserialize and validate the values. |
| *output_values = serialized_values.get<Tensor>(); |
| if (*output_values == nullptr) { |
| return errors::InvalidArgument( |
| "Could not get a tensor from serialized_sparse[", index, ", 1]"); |
| } |
| if (!TensorShapeUtils::IsVector((*output_values)->shape())) { |
| return errors::InvalidArgument( |
| "Expected serialized_sparse[", index, |
| ", 1] to represent a values vector but received shape ", |
| (*output_values)->shape().DebugString()); |
| } |
| if (dtype_ != (*output_values)->dtype()) { |
| return errors::InvalidArgument( |
| "Requested SparseTensor of type ", DataTypeString(dtype_), |
| " but SparseTensor[", index, |
| "].values.dtype() == ", DataTypeString((*output_values)->dtype())); |
| } |
| if (num_entries != (*output_values)->dim_size(0)) { |
| return errors::InvalidArgument( |
| "Expected row counts of SparseTensor[", index, |
| "].indices and SparseTensor[", index, |
| "].values to match but they do not: ", num_entries, " vs. ", |
| (*output_values)->dim_size(0)); |
| } |
| |
| return Status::OK(); |
| } |
| |
| DataType dtype_; |
| }; |
| |
| REGISTER_KERNEL_BUILDER(Name("DeserializeSparse") |
| .Device(DEVICE_CPU) |
| .TypeConstraint<Variant>("Tserialized"), |
| DeserializeSparseOp) |
| |
| } // namespace |
| |
| } // namespace tensorflow |
