-----------------
-| **`Linux CPU`** | **`Linux GPU`** | **`Mac OS CPU`** | **`Windows CPU`** | **`Android`** |
-|-----------------|---------------------|------------------|-------------------|---------------|
-| [](https://ci.tensorflow.org/job/tensorflow-master-cpu) | [](https://ci.tensorflow.org/job/tensorflow-master-linux-gpu) | [](https://ci.tensorflow.org/job/tensorflow-master-mac) | [](https://ci.tensorflow.org/job/tensorflow-master-win-cmake-py) | [](https://ci.tensorflow.org/job/tensorflow-master-android) [  ](https://bintray.com/google/tensorflow/tensorflow/_latestVersion) |
+
+| **`Documentation`** | **`Linux CPU`** | **`Linux GPU`** | **`Mac OS CPU`** | **`Windows CPU`** | **`Android`** |
+|-----------------|---------------------|------------------|-------------------|---------------|---------------|
+| [](https://www.tensorflow.org/api_docs/) | [](https://ci.tensorflow.org/job/tensorflow-master-cpu) | [](https://ci.tensorflow.org/job/tensorflow-master-linux-gpu) | [](https://ci.tensorflow.org/job/tensorflow-master-mac) | [](https://ci.tensorflow.org/job/tensorflow-master-win-cmake-py) | [](https://ci.tensorflow.org/job/tensorflow-master-android) [  ](https://bintray.com/google/tensorflow/tensorflow/_latestVersion)
**TensorFlow** is an open source software library for numerical computation using
data flow graphs. The graph nodes represent mathematical operations, while
networks research. The system is general enough to be applicable in a wide
variety of other domains, as well.
-**If you want to contribute to TensorFlow, be sure to review the [contribution
-guidelines](CONTRIBUTING.md). This project adheres to TensorFlow's
-[code of conduct](CODE_OF_CONDUCT.md). By participating, you are expected to
-uphold this code.**
-
-**We use [GitHub issues](https://github.com/tensorflow/tensorflow/issues) for
-tracking requests and bugs. So please see
-[TensorFlow Discuss](https://groups.google.com/a/tensorflow.org/forum/#!forum/discuss) for general questions
-and discussion, and please direct specific questions to [Stack Overflow](https://stackoverflow.com/questions/tagged/tensorflow).**
-
-The TensorFlow project strives to abide by generally accepted best practices in open-source software development:
-
-[](https://bestpractices.coreinfrastructure.org/projects/1486)
-
## Installation
*See [Installing TensorFlow](https://www.tensorflow.org/get_started/os_setup.html) for instructions on how to install our release binaries or how to build from source.*
>>> sess.close()
```
+## Contribution guidelines
+
+**If you want to contribute to TensorFlow, be sure to review the [contribution
+guidelines](CONTRIBUTING.md). This project adheres to TensorFlow's
+[code of conduct](CODE_OF_CONDUCT.md). By participating, you are expected to
+uphold this code.**
+
+**We use [GitHub issues](https://github.com/tensorflow/tensorflow/issues) for
+tracking requests and bugs. So please see
+[TensorFlow Discuss](https://groups.google.com/a/tensorflow.org/forum/#!forum/discuss) for general questions
+and discussion, and please direct specific questions to [Stack Overflow](https://stackoverflow.com/questions/tagged/tensorflow).**
+
+The TensorFlow project strives to abide by generally accepted best practices in open-source software development:
+
+[](https://bestpractices.coreinfrastructure.org/projects/1486)
+
## For more information
* [TensorFlow Website](https://www.tensorflow.org)
### Known vulnerabilities
-| Type | Versions affected | Reported by | Additional Information |
-|------|:-----------------:|---------------------------------------|
-| out of bounds read| <=1.4 | TenCent Blade Team | [issue report](https://github.com/tensorflow/tensorflow/issues/14959) |
+| Type | Versions affected | Reported by | Additional Information |
+|-------------------|:-----------------:|--------------------|-----------------------------|
+| out of bounds read| <=1.4 | TenCent Blade Team | [issue report](https://github.com/tensorflow/tensorflow/issues/14959) |
fi
# Set all env variables
-"$PYTHON_BIN_PATH" configure.py
+CONFIGURE_DIR=$(dirname "$0")
+"$PYTHON_BIN_PATH" "${CONFIGURE_DIR}/configure.py" "$@"
echo "Configuration finished"
from __future__ import division
from __future__ import print_function
+import argparse
import errno
import os
import platform
from distutils.spawn import find_executable as which
# pylint: enable=g-import-not-at-top
-_TF_BAZELRC = os.path.join(os.path.dirname(os.path.abspath(__file__)),
- '.tf_configure.bazelrc')
-_TF_WORKSPACE = os.path.join(os.path.dirname(os.path.abspath(__file__)),
- 'WORKSPACE')
_DEFAULT_CUDA_VERSION = '9.0'
_DEFAULT_CUDNN_VERSION = '7'
_DEFAULT_CUDA_COMPUTE_CAPABILITIES = '3.5,5.2'
_DEFAULT_PROMPT_ASK_ATTEMPTS = 10
+_TF_WORKSPACE_ROOT = os.path.abspath(os.path.dirname(__file__))
+_TF_BAZELRC_FILENAME = '.tf_configure.bazelrc'
+_TF_BAZELRC = os.path.join(_TF_WORKSPACE_ROOT, _TF_BAZELRC_FILENAME)
+_TF_WORKSPACE = os.path.join(_TF_WORKSPACE_ROOT, 'WORKSPACE')
+
class UserInputError(Exception):
pass
f.write(newdata)
-def remove_line_with(filename, token):
- """Remove lines that contain token from file.
-
- Args:
- filename: string for filename.
- token: string token to check if to remove a line from file or not.
- """
- with open(filename, 'r') as f:
- filedata = f.read()
-
- with open(filename, 'w') as f:
- for line in filedata.strip().split('\n'):
- if token not in line:
- f.write(line + '\n')
-
-
def write_to_bazelrc(line):
with open(_TF_BAZELRC, 'a') as f:
f.write(line + '\n')
environ_cp['PYTHON_BIN_PATH'] = python_bin_path
# Write tools/python_bin_path.sh
- with open('tools/python_bin_path.sh', 'w') as f:
+ with open(os.path.join(
+ _TF_WORKSPACE_ROOT, 'tools', 'python_bin_path.sh'), 'w') as f:
f.write('export PYTHON_BIN_PATH="%s"' % python_bin_path)
-def reset_tf_configure_bazelrc():
+def reset_tf_configure_bazelrc(workspace_path):
"""Reset file that contains customized config settings."""
open(_TF_BAZELRC, 'w').close()
-
- home = os.path.expanduser('~')
- if not os.path.exists('.bazelrc'):
- if os.path.exists(os.path.join(home, '.bazelrc')):
- with open('.bazelrc', 'a') as f:
- f.write('import %s/.bazelrc\n' % home.replace('\\', '/'))
+ bazelrc_path = os.path.join(workspace_path, '.bazelrc')
+
+ data = []
+ if os.path.exists(bazelrc_path):
+ with open(bazelrc_path, 'r') as f:
+ data = f.read().splitlines()
+ with open(bazelrc_path, 'w') as f:
+ for l in data:
+ if _TF_BAZELRC_FILENAME in l:
+ continue
+ f.write('%s\n' % l)
+ if is_windows():
+ tf_bazelrc_path = _TF_BAZELRC.replace("\\", "/")
else:
- open('.bazelrc', 'w').close()
-
- remove_line_with('.bazelrc', 'tf_configure')
- with open('.bazelrc', 'a') as f:
- f.write('import %workspace%/.tf_configure.bazelrc\n')
+ tf_bazelrc_path = _TF_BAZELRC
+ f.write('import %s\n' % tf_bazelrc_path)
def cleanup_makefile():
These files could interfere with Bazel parsing.
"""
- makefile_download_dir = 'tensorflow/contrib/makefile/downloads'
+ makefile_download_dir = os.path.join(
+ _TF_WORKSPACE_ROOT, 'tensorflow', 'contrib', 'makefile', 'downloads')
if os.path.isdir(makefile_download_dir):
for root, _, filenames in os.walk(makefile_download_dir):
for f in filenames:
if which('bazel') is None:
print('Cannot find bazel. Please install bazel.')
sys.exit(0)
- curr_version = run_shell(['bazel', '--batch', 'version'])
+ curr_version = run_shell(['bazel', '--batch', '--bazelrc=/dev/null', 'version'])
for line in curr_version.split('\n'):
if 'Build label: ' in line:
for opt in cc_opt_flags.split():
write_to_bazelrc('build:opt --copt=%s' % opt)
# It should be safe on the same build host.
- write_to_bazelrc('build:opt --host_copt=-march=native')
+ if not is_ppc64le():
+ write_to_bazelrc('build:opt --host_copt=-march=native')
write_to_bazelrc('build:opt --define with_default_optimizations=true')
# TODO(mikecase): Remove these default defines once we are able to get
# TF Lite targets building without them.
environ_cp,
var_name='HOST_C_COMPILER',
var_default=default_c_host_compiler,
- ask_for_var=('Please specify which C compiler should be used as the host'
+ ask_for_var=('Please specify which C compiler should be used as the host '
'C compiler.'),
check_success=os.path.exists,
error_msg='Invalid C compiler path. %s cannot be found.',
def main():
+ parser = argparse.ArgumentParser()
+ parser.add_argument("--workspace",
+ type=str,
+ default=_TF_WORKSPACE_ROOT,
+ help="The absolute path to your active Bazel workspace.")
+ args = parser.parse_args()
+
# Make a copy of os.environ to be clear when functions and getting and setting
# environment variables.
environ_cp = dict(os.environ)
check_bazel_version('0.5.4')
- reset_tf_configure_bazelrc()
+ reset_tf_configure_bazelrc(args.workspace)
cleanup_makefile()
setup_python(environ_cp)
}
REGISTER_GRADIENT_OP("MaxPoolV2", MaxPoolGradV2Helper);
+Status MaxPool3DGradHelper(const Scope& scope, const Operation& op,
+ const std::vector<Output>& grad_inputs,
+ std::vector<Output>* grad_outputs) {
+ std::vector<int32> ksize;
+ std::vector<int32> strides;
+ string padding;
+ string data_format;
+ auto attrs = op.output(0).node()->attrs();
+ TF_RETURN_IF_ERROR(GetNodeAttr(attrs, "ksize", &ksize));
+ TF_RETURN_IF_ERROR(GetNodeAttr(attrs, "strides", &strides));
+ TF_RETURN_IF_ERROR(GetNodeAttr(attrs, "padding", &padding));
+ TF_RETURN_IF_ERROR(GetNodeAttr(attrs, "data_format", &data_format));
+ MaxPool3DGrad::Attrs grad_attrs;
+ auto dx = MaxPool3DGrad(scope, op.input(0), op.output(0), grad_inputs[0],
+ ksize, strides, padding,
+ grad_attrs.DataFormat(data_format));
+ grad_outputs->push_back(dx);
+ return scope.status();
+}
+REGISTER_GRADIENT_OP("MaxPool3D", MaxPool3DGradHelper);
+
+Status AvgPoolGradHelper(const Scope& scope, const Operation& op,
+ const std::vector<Output>& grad_inputs,
+ std::vector<Output>* grad_outputs) {
+ std::vector<int32> ksize;
+ std::vector<int32> strides;
+ string padding;
+ string data_format;
+ auto attrs = op.output(0).node()->attrs();
+ TF_RETURN_IF_ERROR(GetNodeAttr(attrs, "ksize", &ksize));
+ TF_RETURN_IF_ERROR(GetNodeAttr(attrs, "strides", &strides));
+ TF_RETURN_IF_ERROR(GetNodeAttr(attrs, "padding", &padding));
+ TF_RETURN_IF_ERROR(GetNodeAttr(attrs, "data_format", &data_format));
+ internal::AvgPoolGrad::Attrs grad_attrs;
+ auto dx =
+ internal::AvgPoolGrad(scope, Shape(scope, op.input(0)), grad_inputs[0],
+ ksize, strides, padding,
+ grad_attrs.DataFormat(data_format));
+ grad_outputs->push_back(dx);
+ return scope.status();
+}
+REGISTER_GRADIENT_OP("AvgPool", AvgPoolGradHelper);
+
+Status AvgPool3DGradHelper(const Scope& scope, const Operation& op,
+ const std::vector<Output>& grad_inputs,
+ std::vector<Output>* grad_outputs) {
+ std::vector<int32> ksize;
+ std::vector<int32> strides;
+ string padding;
+ string data_format;
+ auto attrs = op.output(0).node()->attrs();
+ TF_RETURN_IF_ERROR(GetNodeAttr(attrs, "ksize", &ksize));
+ TF_RETURN_IF_ERROR(GetNodeAttr(attrs, "strides", &strides));
+ TF_RETURN_IF_ERROR(GetNodeAttr(attrs, "padding", &padding));
+ TF_RETURN_IF_ERROR(GetNodeAttr(attrs, "data_format", &data_format));
+ AvgPool3DGrad::Attrs grad_attrs;
+ auto dx = AvgPool3DGrad(scope, Shape(scope, op.input(0)), grad_inputs[0],
+ ksize, strides, padding,
+ grad_attrs.DataFormat(data_format));
+ grad_outputs->push_back(dx);
+ return scope.status();
+}
+REGISTER_GRADIENT_OP("AvgPool3D", AvgPool3DGradHelper);
+
Status LRNGradHelper(const Scope& scope, const Operation& op,
const std::vector<Output>& grad_inputs,
std::vector<Output>* grad_outputs) {
using ops::L2Loss;
using ops::LogSoftmax;
using ops::LRN;
+using ops::AvgPool;
+using ops::AvgPool3D;
using ops::MaxPool;
using ops::MaxPoolV2;
+using ops::MaxPool3D;
using ops::Placeholder;
using ops::Relu;
using ops::Relu6;
// Sets tensor with random values, ensuring that the max value is largest by
// a reasonable amount.
- // This is an issue for MaxPool and MaxPoolV2, in which perturbations by the
- // numeric gradient computation in the gradient checker can change the max
- // value if values are too close together.
+ // This is an issue for MaxPool, MaxPoolV2 and MaxPool3D, in which
+ // perturbations by the numeric gradient computation in the gradient checker
+ // can change the max value if values are too close together.
template <typename T>
void SetRandomValuesWithBumpedMax(Tensor* tensor) {
auto tensor_flat = tensor->flat<T>();
RunTest(x, x_init_value, y, y_shape);
}
+TEST_F(NNGradTest, MaxPool3DGradHelper) {
+ TensorShape x_shape({1, 3, 3, 3, 1});
+ TensorShape y_shape({1, 1, 1, 1, 1});
+ auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape));
+ // Setup window and strides so that we only do one MaxPool3D.
+ const std::vector<int> ksize{1, 3, 3, 3, 1};
+ const std::vector<int> strides{1, 3, 3, 3, 1};
+ auto y = MaxPool3D(scope_, x, ksize, strides, "VALID");
+ Tensor x_init_value = Tensor(DT_FLOAT, x_shape);
+ SetRandomValuesWithBumpedMax<float>(&x_init_value);
+ RunTest(x, x_init_value, y, y_shape);
+}
+
+TEST_F(NNGradTest, AvgPoolGradHelper) {
+ TensorShape x_shape({1, 2, 2, 1});
+ TensorShape y_shape({1, 1, 1, 1});
+ auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape));
+ // Setup window and strides so that we only do one AvgPool.
+ const std::vector<int> ksize{1, 2, 2, 1};
+ const std::vector<int> strides{1, 2, 2, 1};
+ auto y = AvgPool(scope_, x, ksize, strides, "SAME");
+ RunTest(x, x_shape, y, y_shape);
+}
+
+TEST_F(NNGradTest, AvgPool3DGradHelper) {
+ TensorShape x_shape({1, 3, 3, 3, 1});
+ TensorShape y_shape({1, 1, 1, 1, 1});
+ auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape));
+ // Setup window and strides so that we only do one AvgPool3D.
+ const std::vector<int> ksize{1, 3, 3, 3, 1};
+ const std::vector<int> strides{1, 3, 3, 3, 1};
+ auto y = AvgPool3D(scope_, x, ksize, strides, "SAME");
+ RunTest(x, x_shape, y, y_shape);
+}
+
TEST_F(NNGradTest, LRN){
TensorShape x_shape({1, 1, 2, 1});
auto x = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(x_shape));
/// Adds tracing information `run_meta` to profiler. A `run_meta` is
/// generated by a TensorFlow session run call. `step` is the key
/// to the `run_meta`. When calling ProfileXXX methods, caller can specify
- /// `step` in `options` to seletively profile the corresponding `run_meta`.
+ /// `step` in `options` to selectively profile the corresponding `run_meta`.
/// Multiple different `run_meta` can be keyed by the same `step` in order
/// to group them together.
void AddStep(int64 step, const RunMetadata& run_meta);
/// Profiles the model by organizing nodes in graph structure.
- /// Each node is an op and the nodes are contected by the op inputs/outputs.
+ /// Each node is an op and the nodes are connected by the op inputs/outputs.
GraphNodeProto ProfileGraph(const Options& options);
/// Profiles the model by organizing nodes in name scope structure.
/// Each node is an op, and nodes are organized by the ops' name
- /// scope, similar to a filesystem tree.
+ /// scope, similar to a file system tree.
/// E.g. /foo is the root of operation /foo/matmul_1 and foo/conv_2.
GraphNodeProto ProfileNameScope(const Options& options);
#define __CUDACC__
#include "crt/host_config.h"
-int main(void) { return 0; }
+int main(void) {
+ return 0;
+}
--- /dev/null
+# Copyright 2018 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.
+# ==============================================================================
+"""Experimental methods for tf.feature_column sequence input."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+
+import abc
+import collections
+
+
+from tensorflow.python.feature_column import feature_column as fc
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor_shape
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import check_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import parsing_ops
+from tensorflow.python.ops import sparse_ops
+from tensorflow.python.ops import variable_scope
+
+# TODO(b/73160931): Fix pydoc.
+# pylint: disable=g-doc-args,missing-docstring,protected-access
+# TODO(b/73827486): Support SequenceExample.
+
+
+def sequence_input_layer(
+ features,
+ feature_columns,
+ weight_collections=None,
+ trainable=True,
+ scope=None):
+ """"Builds input layer for sequence input.
+
+ All `feature_columns` must be sequence dense columns with the same
+ `sequence_length`. The output of this method can be fed into sequence
+ networks, such as RNN.
+
+ The output of this method is a 3D `Tensor` of shape `[batch_size, T, D]`.
+ `T` is the maximum sequence length for this batch, which could differ from
+ batch to batch.
+
+ If multiple `feature_columns` are given with `Di` `num_elements` each, their
+ outputs are concatenated. So, the final `Tensor` has shape
+ `[batch_size, T, D0 + D1 + ... + Dn]`.
+
+ Example:
+
+ ```python
+ rating = sequence_numeric_column('rating')
+ watches = sequence_categorical_column_with_identity(
+ 'watches', num_buckets=1000)
+ watches_embedding = embedding_column(watches, dimension=10)
+ columns = [rating, watches]
+
+ features = tf.parse_example(..., features=make_parse_example_spec(columns))
+ input_layer, sequence_length = sequence_input_layer(features, columns)
+
+ rnn_cell = tf.nn.rnn_cell.BasicRNNCell(hidden_size)
+ outputs, state = tf.nn.dynamic_rnn(
+ rnn_cell, inputs=input_layer, sequence_length=sequence_length)
+ ```
+
+ Returns:
+ An `(input_layer, sequence_length)` tuple where:
+ - input_layer: A float `Tensor` of shape `[batch_size, T, D]`.
+ `T` is the maximum sequence length for this batch, which could differ
+ from batch to batch. `D` is the sum of `num_elements` for all
+ `feature_columns`.
+ - sequence_length: An int `Tensor` of shape `[batch_size]`. The sequence
+ length for each example.
+ Raises:
+ ValueError: If any of the `feature_columns` is the wrong type.
+ """
+ feature_columns = fc._clean_feature_columns(feature_columns)
+ for c in feature_columns:
+ if not isinstance(c, _SequenceDenseColumn):
+ raise ValueError(
+ 'All feature_columns must be of type _SequenceDenseColumn. '
+ 'Given (type {}): {}'.format(type(c), c))
+
+ with variable_scope.variable_scope(
+ scope, default_name='sequence_input_layer', values=features.values()):
+ builder = fc._LazyBuilder(features)
+ output_tensors = []
+ sequence_lengths = []
+ ordered_columns = []
+ for column in sorted(feature_columns, key=lambda x: x.name):
+ ordered_columns.append(column)
+ with variable_scope.variable_scope(
+ None, default_name=column._var_scope_name):
+ dense_tensor, sequence_length = column._get_sequence_dense_tensor(
+ builder,
+ weight_collections=weight_collections,
+ trainable=trainable)
+ # Flattens the final dimension to produce a 3D Tensor.
+ num_elements = column._variable_shape.num_elements()
+ shape = array_ops.shape(dense_tensor)
+ output_tensors.append(
+ array_ops.reshape(
+ dense_tensor,
+ shape=array_ops.concat([shape[:2], [num_elements]], axis=0)))
+ sequence_lengths.append(sequence_length)
+ fc._verify_static_batch_size_equality(output_tensors, ordered_columns)
+ # TODO(b/73160931): Verify sequence_length equality.
+ return array_ops.concat(output_tensors, -1), sequence_lengths[0]
+
+
+# TODO(b/73160931): Add remaining categorical columns.
+def sequence_categorical_column_with_identity(
+ key, num_buckets, default_value=None):
+ return _SequenceCategoricalColumn(
+ fc.categorical_column_with_identity(
+ key=key,
+ num_buckets=num_buckets,
+ default_value=default_value))
+
+
+# TODO(b/73160931): Merge with embedding_column
+def _sequence_embedding_column(
+ categorical_column, dimension, initializer=None, ckpt_to_load_from=None,
+ tensor_name_in_ckpt=None, max_norm=None, trainable=True):
+ if not isinstance(categorical_column, _SequenceCategoricalColumn):
+ raise ValueError(
+ 'categorical_column must be of type _SequenceCategoricalColumn. '
+ 'Given (type {}): {}'.format(
+ type(categorical_column), categorical_column))
+ return _SequenceEmbeddingColumn(
+ fc.embedding_column(
+ categorical_column,
+ dimension=dimension,
+ initializer=initializer,
+ ckpt_to_load_from=ckpt_to_load_from,
+ tensor_name_in_ckpt=tensor_name_in_ckpt,
+ max_norm=max_norm,
+ trainable=trainable))
+
+
+def sequence_numeric_column(
+ key,
+ shape=(1,),
+ default_value=0.,
+ dtype=dtypes.float32):
+ # TODO(b/73160931): Add validations.
+ return _SequenceNumericColumn(
+ key,
+ shape=shape,
+ default_value=default_value,
+ dtype=dtype)
+
+
+class _SequenceDenseColumn(fc._FeatureColumn):
+ """Represents dense sequence data."""
+
+ __metaclass__ = abc.ABCMeta
+
+ TensorSequenceLengthPair = collections.namedtuple( # pylint: disable=invalid-name
+ 'TensorSequenceLengthPair', ['dense_tensor', 'sequence_length'])
+
+ @abc.abstractproperty
+ def _variable_shape(self):
+ """`TensorShape` without batch and sequence dimensions."""
+ pass
+
+ @abc.abstractmethod
+ def _get_sequence_dense_tensor(
+ self, inputs, weight_collections=None, trainable=None):
+ """Returns a `TensorSequenceLengthPair`."""
+ pass
+
+
+def _sequence_length_from_sparse_tensor(sp_tensor, num_elements=1):
+ with ops.name_scope(None, 'sequence_length') as name_scope:
+ row_ids = sp_tensor.indices[:, 0]
+ column_ids = sp_tensor.indices[:, 1]
+ column_ids += array_ops.ones_like(column_ids)
+ seq_length = (
+ math_ops.segment_max(column_ids, segment_ids=row_ids) / num_elements)
+ # If the last n rows do not have ids, seq_length will have shape
+ # [batch_size - n]. Pad the remaining values with zeros.
+ n_pad = array_ops.shape(sp_tensor)[:1] - array_ops.shape(seq_length)[:1]
+ padding = array_ops.zeros(n_pad, dtype=seq_length.dtype)
+ return array_ops.concat([seq_length, padding], axis=0, name=name_scope)
+
+
+class _SequenceCategoricalColumn(
+ fc._CategoricalColumn,
+ collections.namedtuple(
+ '_SequenceCategoricalColumn', ['categorical_column'])):
+
+ @property
+ def name(self):
+ return self.categorical_column.name
+
+ @property
+ def _parse_example_spec(self):
+ return self.categorical_column._parse_example_spec
+
+ def _transform_feature(self, inputs):
+ return self.categorical_column._transform_feature(inputs)
+
+ @property
+ def _num_buckets(self):
+ return self.categorical_column._num_buckets
+
+ def _get_sparse_tensors(self, inputs, weight_collections=None,
+ trainable=None):
+ sparse_tensors = self.categorical_column._get_sparse_tensors(inputs)
+ id_tensor = sparse_tensors.id_tensor
+ weight_tensor = sparse_tensors.weight_tensor
+ # Expands final dimension, so that embeddings are not combined during
+ # embedding lookup.
+ check_id_rank = check_ops.assert_equal(
+ array_ops.rank(id_tensor), 2,
+ data=[
+ 'Column {} expected ID tensor of rank 2. '.format(self.name),
+ 'id_tensor shape: ', array_ops.shape(id_tensor)])
+ with ops.control_dependencies([check_id_rank]):
+ id_tensor = sparse_ops.sparse_reshape(
+ id_tensor,
+ shape=array_ops.concat([id_tensor.dense_shape, [1]], axis=0))
+ if weight_tensor is not None:
+ check_weight_rank = check_ops.assert_equal(
+ array_ops.rank(weight_tensor), 2,
+ data=[
+ 'Column {} expected weight tensor of rank 2.'.format(self.name),
+ 'weight_tensor shape:', array_ops.shape(weight_tensor)])
+ with ops.control_dependencies([check_weight_rank]):
+ weight_tensor = sparse_ops.sparse_reshape(
+ weight_tensor,
+ shape=array_ops.concat([weight_tensor.dense_shape, [1]], axis=0))
+ return fc._CategoricalColumn.IdWeightPair(id_tensor, weight_tensor)
+
+ def _sequence_length(self, inputs):
+ sparse_tensors = self.categorical_column._get_sparse_tensors(inputs)
+ return _sequence_length_from_sparse_tensor(sparse_tensors.id_tensor)
+
+
+class _SequenceEmbeddingColumn(
+ _SequenceDenseColumn,
+ collections.namedtuple('_SequenceEmbeddingColumn', ['embedding_column'])):
+
+ @property
+ def name(self):
+ return self.embedding_column.name
+
+ @property
+ def _parse_example_spec(self):
+ return self.embedding_column._parse_example_spec
+
+ def _transform_feature(self, inputs):
+ return self.embedding_column._transform_feature(inputs)
+
+ @property
+ def _variable_shape(self):
+ return self.embedding_column._variable_shape
+
+ def _get_sequence_dense_tensor(
+ self, inputs, weight_collections=None, trainable=None):
+ dense_tensor = self.embedding_column._get_dense_tensor(
+ inputs=inputs,
+ weight_collections=weight_collections,
+ trainable=trainable)
+ sequence_length = self.embedding_column.categorical_column._sequence_length(
+ inputs)
+ return _SequenceDenseColumn.TensorSequenceLengthPair(
+ dense_tensor=dense_tensor, sequence_length=sequence_length)
+
+
+class _SequenceNumericColumn(
+ _SequenceDenseColumn,
+ collections.namedtuple(
+ '_SequenceNumericColumn',
+ ['key', 'shape', 'default_value', 'dtype'])):
+
+ @property
+ def name(self):
+ return self.key
+
+ @property
+ def _parse_example_spec(self):
+ return {self.key: parsing_ops.VarLenFeature(self.dtype)}
+
+ def _transform_feature(self, inputs):
+ return inputs.get(self.key)
+
+ @property
+ def _variable_shape(self):
+ return tensor_shape.TensorShape(self.shape)
+
+ def _get_sequence_dense_tensor(
+ self, inputs, weight_collections=None, trainable=None):
+ # Do nothing with weight_collections and trainable since no variables are
+ # created in this function.
+ del weight_collections
+ del trainable
+ sp_tensor = inputs.get(self)
+ dense_tensor = sparse_ops.sparse_tensor_to_dense(
+ sp_tensor, default_value=self.default_value)
+ # Reshape into [batch_size, T, variable_shape].
+ dense_shape = array_ops.concat(
+ [array_ops.shape(dense_tensor)[:1], [-1], self._variable_shape],
+ axis=0)
+ dense_tensor = array_ops.reshape(dense_tensor, shape=dense_shape)
+ sequence_length = _sequence_length_from_sparse_tensor(
+ sp_tensor, num_elements=self._variable_shape.num_elements())
+ return _SequenceDenseColumn.TensorSequenceLengthPair(
+ dense_tensor=dense_tensor, sequence_length=sequence_length)
+
+# pylint: enable=g-doc-args,missing-docstring,protected-access
--- /dev/null
+# Copyright 2018 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.
+# ==============================================================================
+"""Tests for sequential_feature_column."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import numpy as np
+
+from tensorflow.contrib.feature_column.python.feature_column import sequential_feature_column as sfc
+from tensorflow.python.feature_column.feature_column import _LazyBuilder
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import sparse_tensor
+from tensorflow.python.platform import test
+from tensorflow.python.training import monitored_session
+
+
+class SequenceInputLayerTest(test.TestCase):
+
+ def test_embedding_column(self):
+ vocabulary_size = 3
+ sparse_input_a = sparse_tensor.SparseTensorValue(
+ # example 0, ids [2]
+ # example 1, ids [0, 1]
+ indices=((0, 0), (1, 0), (1, 1)),
+ values=(2, 0, 1),
+ dense_shape=(2, 2))
+ sparse_input_b = sparse_tensor.SparseTensorValue(
+ # example 0, ids [1]
+ # example 1, ids [2, 0]
+ indices=((0, 0), (1, 0), (1, 1)),
+ values=(1, 2, 0),
+ dense_shape=(2, 2))
+
+ embedding_dimension_a = 2
+ embedding_values_a = (
+ (1., 2.), # id 0
+ (3., 4.), # id 1
+ (5., 6.) # id 2
+ )
+ embedding_dimension_b = 3
+ embedding_values_b = (
+ (11., 12., 13.), # id 0
+ (14., 15., 16.), # id 1
+ (17., 18., 19.) # id 2
+ )
+ def _get_initializer(embedding_dimension, embedding_values):
+ def _initializer(shape, dtype, partition_info):
+ self.assertAllEqual((vocabulary_size, embedding_dimension), shape)
+ self.assertEqual(dtypes.float32, dtype)
+ self.assertIsNone(partition_info)
+ return embedding_values
+ return _initializer
+
+ expected_input_layer = [
+ # example 0, ids_a [2], ids_b [1]
+ [[5., 6., 14., 15., 16.], [0., 0., 0., 0., 0.]],
+ # example 1, ids_a [0, 1], ids_b [2, 0]
+ [[1., 2., 17., 18., 19.], [3., 4., 11., 12., 13.]],
+ ]
+ expected_sequence_length = [1, 2]
+
+ categorical_column_a = sfc.sequence_categorical_column_with_identity(
+ key='aaa', num_buckets=vocabulary_size)
+ embedding_column_a = sfc._sequence_embedding_column(
+ categorical_column_a, dimension=embedding_dimension_a,
+ initializer=_get_initializer(embedding_dimension_a, embedding_values_a))
+ categorical_column_b = sfc.sequence_categorical_column_with_identity(
+ key='bbb', num_buckets=vocabulary_size)
+ embedding_column_b = sfc._sequence_embedding_column(
+ categorical_column_b, dimension=embedding_dimension_b,
+ initializer=_get_initializer(embedding_dimension_b, embedding_values_b))
+
+ input_layer, sequence_length = sfc.sequence_input_layer(
+ features={
+ 'aaa': sparse_input_a,
+ 'bbb': sparse_input_b,
+ },
+ # Test that columns are reordered alphabetically.
+ feature_columns=[embedding_column_b, embedding_column_a])
+
+ global_vars = ops.get_collection(ops.GraphKeys.GLOBAL_VARIABLES)
+ self.assertItemsEqual(
+ ('sequence_input_layer/aaa_embedding/embedding_weights:0',
+ 'sequence_input_layer/bbb_embedding/embedding_weights:0'),
+ tuple([v.name for v in global_vars]))
+ with monitored_session.MonitoredSession() as sess:
+ self.assertAllEqual(embedding_values_a, global_vars[0].eval(session=sess))
+ self.assertAllEqual(embedding_values_b, global_vars[1].eval(session=sess))
+ self.assertAllEqual(expected_input_layer, input_layer.eval(session=sess))
+ self.assertAllEqual(
+ expected_sequence_length, sequence_length.eval(session=sess))
+
+ def test_numeric_column(self):
+ sparse_input = sparse_tensor.SparseTensorValue(
+ # example 0, values [[0.], [1]]
+ # example 1, [[10.]]
+ indices=((0, 0), (0, 1), (1, 0)),
+ values=(0., 1., 10.),
+ dense_shape=(2, 2))
+ expected_input_layer = [
+ [[0.], [1.]],
+ [[10.], [0.]],
+ ]
+ expected_sequence_length = [2, 1]
+ numeric_column = sfc.sequence_numeric_column('aaa')
+
+ input_layer, sequence_length = sfc.sequence_input_layer(
+ features={'aaa': sparse_input},
+ feature_columns=[numeric_column])
+
+ with monitored_session.MonitoredSession() as sess:
+ self.assertAllEqual(expected_input_layer, input_layer.eval(session=sess))
+ self.assertAllEqual(
+ expected_sequence_length, sequence_length.eval(session=sess))
+
+ def test_numeric_column_multi_dim(self):
+ """Tests sequence_input_layer for multi-dimensional numeric_column."""
+ sparse_input = sparse_tensor.SparseTensorValue(
+ # example 0, values [[[0., 1.], [2., 3.]], [[4., 5.], [6., 7.]]]
+ # example 1, [[[10., 11.], [12., 13.]]]
+ indices=((0, 0), (0, 1), (0, 2), (0, 3), (0, 4), (0, 5), (0, 6), (0, 7),
+ (1, 0), (1, 1), (1, 2), (1, 3)),
+ values=(0., 1., 2., 3., 4., 5., 6., 7., 10., 11., 12., 13.),
+ dense_shape=(2, 8))
+ # The output of numeric_column._get_dense_tensor should be flattened.
+ expected_input_layer = [
+ [[0., 1., 2., 3.], [4., 5., 6., 7.]],
+ [[10., 11., 12., 13.], [0., 0., 0., 0.]],
+ ]
+ expected_sequence_length = [2, 1]
+ numeric_column = sfc.sequence_numeric_column('aaa', shape=(2, 2))
+
+ input_layer, sequence_length = sfc.sequence_input_layer(
+ features={'aaa': sparse_input},
+ feature_columns=[numeric_column])
+
+ with monitored_session.MonitoredSession() as sess:
+ self.assertAllEqual(expected_input_layer, input_layer.eval(session=sess))
+ self.assertAllEqual(
+ expected_sequence_length, sequence_length.eval(session=sess))
+
+
+def _assert_sparse_tensor_value(test_case, expected, actual):
+ test_case.assertEqual(np.int64, np.array(actual.indices).dtype)
+ test_case.assertAllEqual(expected.indices, actual.indices)
+
+ test_case.assertEqual(
+ np.array(expected.values).dtype, np.array(actual.values).dtype)
+ test_case.assertAllEqual(expected.values, actual.values)
+
+ test_case.assertEqual(np.int64, np.array(actual.dense_shape).dtype)
+ test_case.assertAllEqual(expected.dense_shape, actual.dense_shape)
+
+
+class SequenceCategoricalColumnWithIdentityTest(test.TestCase):
+
+ def test_get_sparse_tensors(self):
+ column = sfc.sequence_categorical_column_with_identity(
+ 'aaa', num_buckets=3)
+ inputs = sparse_tensor.SparseTensorValue(
+ indices=((0, 0), (1, 0), (1, 1)),
+ values=(1, 2, 0),
+ dense_shape=(2, 2))
+ expected_sparse_ids = sparse_tensor.SparseTensorValue(
+ indices=((0, 0, 0), (1, 0, 0), (1, 1, 0)),
+ values=np.array((1, 2, 0), dtype=np.int64),
+ dense_shape=(2, 2, 1))
+
+ id_weight_pair = column._get_sparse_tensors(_LazyBuilder({'aaa': inputs}))
+
+ self.assertIsNone(id_weight_pair.weight_tensor)
+ with monitored_session.MonitoredSession() as sess:
+ _assert_sparse_tensor_value(
+ self,
+ expected_sparse_ids,
+ id_weight_pair.id_tensor.eval(session=sess))
+
+ def test_get_sparse_tensors_inputs3d(self):
+ """Tests _get_sparse_tensors when the input is already 3D Tensor."""
+ column = sfc.sequence_categorical_column_with_identity(
+ 'aaa', num_buckets=3)
+ inputs = sparse_tensor.SparseTensorValue(
+ indices=((0, 0, 0), (1, 0, 0), (1, 1, 0)),
+ values=(1, 2, 0),
+ dense_shape=(2, 2, 1))
+
+ with self.assertRaisesRegexp(
+ errors.InvalidArgumentError,
+ r'Column aaa expected ID tensor of rank 2\.\s*'
+ r'id_tensor shape:\s*\[2 2 1\]'):
+ id_weight_pair = column._get_sparse_tensors(
+ _LazyBuilder({'aaa': inputs}))
+ with monitored_session.MonitoredSession() as sess:
+ id_weight_pair.id_tensor.eval(session=sess)
+
+ def test_sequence_length(self):
+ column = sfc.sequence_categorical_column_with_identity(
+ 'aaa', num_buckets=3)
+ inputs = sparse_tensor.SparseTensorValue(
+ indices=((0, 0), (1, 0), (1, 1)),
+ values=(1, 2, 0),
+ dense_shape=(2, 2))
+ expected_sequence_length = [1, 2]
+
+ sequence_length = column._sequence_length(_LazyBuilder({'aaa': inputs}))
+
+ with monitored_session.MonitoredSession() as sess:
+ self.assertAllEqual(
+ expected_sequence_length, sequence_length.eval(session=sess))
+
+ def test_sequence_length_with_zeros(self):
+ column = sfc.sequence_categorical_column_with_identity(
+ 'aaa', num_buckets=3)
+ inputs = sparse_tensor.SparseTensorValue(
+ indices=((1, 0), (3, 0), (3, 1)),
+ values=(1, 2, 0),
+ dense_shape=(5, 2))
+ expected_sequence_length = [0, 1, 0, 2, 0]
+
+ sequence_length = column._sequence_length(_LazyBuilder({'aaa': inputs}))
+
+ with monitored_session.MonitoredSession() as sess:
+ self.assertAllEqual(
+ expected_sequence_length, sequence_length.eval(session=sess))
+
+
+class SequenceEmbeddingColumnTest(test.TestCase):
+
+ def test_get_sequence_dense_tensor(self):
+ vocabulary_size = 3
+ sparse_input = sparse_tensor.SparseTensorValue(
+ # example 0, ids [2]
+ # example 1, ids [0, 1]
+ # example 2, ids []
+ # example 3, ids [1]
+ indices=((0, 0), (1, 0), (1, 1), (3, 0)),
+ values=(2, 0, 1, 1),
+ dense_shape=(4, 2))
+
+ embedding_dimension = 2
+ embedding_values = (
+ (1., 2.), # id 0
+ (3., 5.), # id 1
+ (7., 11.) # id 2
+ )
+ def _initializer(shape, dtype, partition_info):
+ self.assertAllEqual((vocabulary_size, embedding_dimension), shape)
+ self.assertEqual(dtypes.float32, dtype)
+ self.assertIsNone(partition_info)
+ return embedding_values
+
+ expected_lookups = [
+ # example 0, ids [2]
+ [[7., 11.], [0., 0.]],
+ # example 1, ids [0, 1]
+ [[1., 2.], [3., 5.]],
+ # example 2, ids []
+ [[0., 0.], [0., 0.]],
+ # example 3, ids [1]
+ [[3., 5.], [0., 0.]],
+ ]
+
+ categorical_column = sfc.sequence_categorical_column_with_identity(
+ key='aaa', num_buckets=vocabulary_size)
+ embedding_column = sfc._sequence_embedding_column(
+ categorical_column, dimension=embedding_dimension,
+ initializer=_initializer)
+
+ embedding_lookup, _ = embedding_column._get_sequence_dense_tensor(
+ _LazyBuilder({'aaa': sparse_input}))
+
+ global_vars = ops.get_collection(ops.GraphKeys.GLOBAL_VARIABLES)
+ self.assertItemsEqual(
+ ('embedding_weights:0',), tuple([v.name for v in global_vars]))
+ with monitored_session.MonitoredSession() as sess:
+ self.assertAllEqual(embedding_values, global_vars[0].eval(session=sess))
+ self.assertAllEqual(expected_lookups, embedding_lookup.eval(session=sess))
+
+ def test_sequence_length(self):
+ vocabulary_size = 3
+ sparse_input = sparse_tensor.SparseTensorValue(
+ # example 0, ids [2]
+ # example 1, ids [0, 1]
+ indices=((0, 0), (1, 0), (1, 1)),
+ values=(2, 0, 1),
+ dense_shape=(2, 2))
+ expected_sequence_length = [1, 2]
+
+ categorical_column = sfc.sequence_categorical_column_with_identity(
+ key='aaa', num_buckets=vocabulary_size)
+ embedding_column = sfc._sequence_embedding_column(
+ categorical_column, dimension=2)
+
+ _, sequence_length = embedding_column._get_sequence_dense_tensor(
+ _LazyBuilder({'aaa': sparse_input}))
+
+ with monitored_session.MonitoredSession() as sess:
+ self.assertAllEqual(
+ expected_sequence_length, sequence_length.eval(session=sess))
+
+ def test_sequence_length_with_empty_rows(self):
+ """Tests _sequence_length when some examples do not have ids."""
+ vocabulary_size = 3
+ sparse_input = sparse_tensor.SparseTensorValue(
+ # example 0, ids []
+ # example 1, ids [2]
+ # example 2, ids [0, 1]
+ # example 3, ids []
+ # example 4, ids [1]
+ # example 5, ids []
+ indices=((1, 0), (2, 0), (2, 1), (4, 0)),
+ values=(2, 0, 1, 1),
+ dense_shape=(6, 2))
+ expected_sequence_length = [0, 1, 2, 0, 1, 0]
+
+ categorical_column = sfc.sequence_categorical_column_with_identity(
+ key='aaa', num_buckets=vocabulary_size)
+ embedding_column = sfc._sequence_embedding_column(
+ categorical_column, dimension=2)
+
+ _, sequence_length = embedding_column._get_sequence_dense_tensor(
+ _LazyBuilder({'aaa': sparse_input}))
+
+ with monitored_session.MonitoredSession() as sess:
+ self.assertAllEqual(
+ expected_sequence_length, sequence_length.eval(session=sess))
+
+
+class SequenceNumericColumnTest(test.TestCase):
+
+ def test_get_sequence_dense_tensor(self):
+ sparse_input = sparse_tensor.SparseTensorValue(
+ # example 0, values [[0.], [1]]
+ # example 1, [[10.]]
+ indices=((0, 0), (0, 1), (1, 0)),
+ values=(0., 1., 10.),
+ dense_shape=(2, 2))
+ expected_dense_tensor = [
+ [[0.], [1.]],
+ [[10.], [0.]],
+ ]
+ numeric_column = sfc.sequence_numeric_column('aaa')
+
+ dense_tensor, _ = numeric_column._get_sequence_dense_tensor(
+ _LazyBuilder({'aaa': sparse_input}))
+
+ with monitored_session.MonitoredSession() as sess:
+ self.assertAllEqual(
+ expected_dense_tensor, dense_tensor.eval(session=sess))
+
+ def test_get_sequence_dense_tensor_with_shape(self):
+ """Tests get_sequence_dense_tensor with shape !=(1,)."""
+ sparse_input = sparse_tensor.SparseTensorValue(
+ # example 0, values [[0., 1., 2.], [3., 4., 5.]]
+ # example 1, [[10., 11., 12.]]
+ indices=((0, 0), (0, 1), (0, 2), (0, 3), (0, 4), (0, 5),
+ (1, 0), (1, 1), (1, 2)),
+ values=(0., 1., 2., 3., 4., 5., 10., 11., 12.),
+ dense_shape=(2, 6))
+ expected_dense_tensor = [
+ [[0., 1., 2.], [3., 4., 5.]],
+ [[10., 11., 12.], [0., 0., 0.]],
+ ]
+ numeric_column = sfc.sequence_numeric_column('aaa', shape=(3,))
+
+ dense_tensor, _ = numeric_column._get_sequence_dense_tensor(
+ _LazyBuilder({'aaa': sparse_input}))
+
+ with monitored_session.MonitoredSession() as sess:
+ self.assertAllEqual(
+ expected_dense_tensor, dense_tensor.eval(session=sess))
+
+ def test_get_dense_tensor_multi_dim(self):
+ """Tests get_sequence_dense_tensor for multi-dim numeric_column."""
+ sparse_input = sparse_tensor.SparseTensorValue(
+ # example 0, values [[[0., 1.], [2., 3.]], [[4., 5.], [6., 7.]]]
+ # example 1, [[[10., 11.], [12., 13.]]]
+ indices=((0, 0), (0, 1), (0, 2), (0, 3), (0, 4), (0, 5), (0, 6), (0, 7),
+ (1, 0), (1, 1), (1, 2), (1, 3)),
+ values=(0., 1., 2., 3., 4., 5., 6., 7., 10., 11., 12., 13.),
+ dense_shape=(2, 8))
+ expected_dense_tensor = [
+ [[[0., 1.], [2., 3.]], [[4., 5.], [6., 7.]]],
+ [[[10., 11.], [12., 13.]], [[0., 0.], [0., 0.]]],
+ ]
+ numeric_column = sfc.sequence_numeric_column('aaa', shape=(2, 2))
+
+ dense_tensor, _ = numeric_column._get_sequence_dense_tensor(
+ _LazyBuilder({'aaa': sparse_input}))
+
+ with monitored_session.MonitoredSession() as sess:
+ self.assertAllEqual(
+ expected_dense_tensor, dense_tensor.eval(session=sess))
+
+ def test_sequence_length(self):
+ sparse_input = sparse_tensor.SparseTensorValue(
+ # example 0, values [[0., 1., 2.], [3., 4., 5.]]
+ # example 1, [[10., 11., 12.]]
+ indices=((0, 0), (0, 1), (0, 2), (0, 3), (0, 4), (0, 5),
+ (1, 0), (1, 1), (1, 2)),
+ values=(0., 1., 2., 3., 4., 5., 10., 11., 12.),
+ dense_shape=(2, 6))
+ expected_sequence_length = [2, 1]
+ numeric_column = sfc.sequence_numeric_column('aaa', shape=(3,))
+
+ _, sequence_length = numeric_column._get_sequence_dense_tensor(
+ _LazyBuilder({'aaa': sparse_input}))
+
+ with monitored_session.MonitoredSession() as sess:
+ self.assertAllEqual(
+ expected_sequence_length, sequence_length.eval(session=sess))
+
+ def test_sequence_length_with_shape(self):
+ """Tests _sequence_length with shape !=(1,)."""
+ sparse_input = sparse_tensor.SparseTensorValue(
+ # example 0, values [[0.], [1]]
+ # example 1, [[10.]]
+ indices=((0, 0), (0, 1), (1, 0)),
+ values=(0., 1., 10.),
+ dense_shape=(2, 2))
+ expected_sequence_length = [2, 1]
+ numeric_column = sfc.sequence_numeric_column('aaa')
+
+ _, sequence_length = numeric_column._get_sequence_dense_tensor(
+ _LazyBuilder({'aaa': sparse_input}))
+
+ with monitored_session.MonitoredSession() as sess:
+ self.assertAllEqual(
+ expected_sequence_length, sequence_length.eval(session=sess))
+
+ def test_sequence_length_with_empty_rows(self):
+ """Tests _sequence_length when some examples do not have ids."""
+ sparse_input = sparse_tensor.SparseTensorValue(
+ # example 0, values []
+ # example 1, values [[0.], [1.]]
+ # example 2, [[2.]]
+ # example 3, values []
+ # example 4, [[3.]]
+ # example 5, values []
+ indices=((1, 0), (1, 1), (2, 0), (4, 0)),
+ values=(0., 1., 2., 3.),
+ dense_shape=(6, 2))
+ expected_sequence_length = [0, 2, 1, 0, 1, 0]
+ numeric_column = sfc.sequence_numeric_column('aaa')
+
+ _, sequence_length = numeric_column._get_sequence_dense_tensor(
+ _LazyBuilder({'aaa': sparse_input}))
+
+ with monitored_session.MonitoredSession() as sess:
+ self.assertAllEqual(
+ expected_sequence_length, sequence_length.eval(session=sess))
+
+
+if __name__ == '__main__':
+ test.main()
class SummariesTest(test.TestCase):
- def _test_add_gan_model_image_summaries_impl(
- self, get_model_fn, expected_num_summary_ops, model_summaries):
- summaries.add_gan_model_image_summaries(
- get_model_fn(), grid_size=2, model_summaries=model_summaries)
+ def _test_add_gan_model_image_summaries_impl(self, get_model_fn,
+ expected_num_summary_ops,
+ model_summaries):
+ summaries.add_gan_model_image_summaries(get_model_fn(), grid_size=2,
+ model_summaries=model_summaries)
self.assertEquals(expected_num_summary_ops,
len(ops.get_collection(ops.GraphKeys.SUMMARIES)))
'avg_pool2d', 'avg_pool3d', 'batch_norm', 'bias_add', 'conv2d', 'conv3d',
'conv2d_in_plane', 'conv2d_transpose', 'conv3d_transpose', 'convolution',
'convolution2d', 'convolution2d_in_plane', 'convolution2d_transpose',
- 'convolution3d', 'convolution3d_transpose', 'dense_to_sparse', 'dropout',
- 'elu', 'flatten', 'fully_connected', 'GDN', 'gdn', 'images_to_sequence',
- 'layer_norm', 'linear', 'pool', 'max_pool2d', 'max_pool3d',
- 'one_hot_encoding', 'relu', 'relu6', 'repeat', 'scale_gradient',
- 'separable_conv2d', 'separable_convolution2d', 'sequence_to_images',
- 'softmax', 'spatial_softmax', 'stack', 'unit_norm',
+ 'convolution3d', 'convolution3d_transpose', 'dense_to_sparse',
+ 'dropout', 'elu', 'flatten', 'fully_connected', 'GDN', 'gdn',
+ 'images_to_sequence', 'layer_norm', 'linear', 'pool', 'max_pool2d',
+ 'max_pool3d', 'one_hot_encoding', 'relu', 'relu6', 'repeat',
+ 'scale_gradient', 'separable_conv2d', 'separable_convolution2d',
+ 'sequence_to_images', 'softmax', 'spatial_softmax', 'stack', 'unit_norm',
'legacy_fully_connected', 'legacy_linear', 'legacy_relu', 'maxout'
]
num_batches = -1
else:
num_batches = num_batches // height
- reshaped = array_ops.reshape(inputs, [width, num_batches, height, depth])
+ reshaped = array_ops.reshape(inputs,
+ [width, num_batches, height, depth])
if output_data_format == 'channels_first':
outputs = array_ops.transpose(reshaped, [1, 3, 2, 0])
else:
num_time_steps = 11
num_channels = 5
desired_height = 7
- sequence = np.random.uniform(
- size=(num_time_steps, num_batches, num_channels)).astype(np.float32)
+ sequence = np.random.uniform(size=(num_time_steps,
+ num_batches,
+ num_channels)).astype(np.float32)
output = _layers.sequence_to_images(sequence, desired_height)
self.assertListEqual(output.get_shape().as_list(), [2, 7, 11, 5])
num_time_steps = 11
num_channels = 5
desired_height = 7
- sequence = np.random.uniform(
- size=(num_time_steps, num_batches, num_channels)).astype(np.float32)
- output = _layers.sequence_to_images(
- sequence, desired_height, output_data_format='channels_first')
+ sequence = np.random.uniform(size=(num_time_steps,
+ num_batches,
+ num_channels)).astype(np.float32)
+ output = _layers.sequence_to_images(sequence,
+ desired_height,
+ output_data_format='channels_first')
self.assertListEqual(output.get_shape().as_list(), [2, 5, 7, 11])
import android.graphics.Bitmap;
import android.os.SystemClock;
import android.util.Log;
+
+import org.tensorflow.lite.Interpreter;
+
import java.io.BufferedReader;
import java.io.FileInputStream;
import java.io.IOException;
import java.util.List;
import java.util.Map;
import java.util.PriorityQueue;
-import org.tensorflow.lite.Interpreter;
-/** Classifies images with Tensorflow Lite. */
+/**
+ * Classifies images with Tensorflow Lite.
+ */
public abstract class ImageClassifier {
/** Tag for the {@link Log}. */
package com.example.android.tflitecamerademo;
import android.app.Activity;
+
import java.io.IOException;
/**
- * This classifier works with the Inception-v3 slim model. It applies floating point inference
- * rather than using a quantized model.
+ * This classifier works with the Inception-v3 slim model.
+ * It applies floating point inference rather than using a quantized model.
*/
public class ImageClassifierFloatInception extends ImageClassifier {
- /** The inception net requires additional normalization of the used input. */
+ /**
+ * The inception net requires additional normalization of the used input.
+ */
private static final int IMAGE_MEAN = 128;
-
private static final float IMAGE_STD = 128.0f;
/**
- * An array to hold inference results, to be feed into Tensorflow Lite as outputs. This isn't part
- * of the super class, because we need a primitive array here.
+ * An array to hold inference results, to be feed into Tensorflow Lite as outputs.
+ * This isn't part of the super class, because we need a primitive array here.
*/
private float[][] labelProbArray = null;
package com.example.android.tflitecamerademo;
import android.app.Activity;
+
import java.io.IOException;
-/** This classifier works with the quantized MobileNet model. */
+/**
+ * This classifier works with the quantized MobileNet model.
+ */
public class ImageClassifierQuantizedMobileNet extends ImageClassifier {
/**
- * An array to hold inference results, to be feed into Tensorflow Lite as outputs. This isn't part
- * of the super class, because we need a primitive array here.
+ * An array to hold inference results, to be feed into Tensorflow Lite as outputs.
+ * This isn't part of the super class, because we need a primitive array here.
*/
private byte[][] labelProbArray = null;
#include <string.h>
#include "tensorflow/contrib/lite/builtin_op_data.h"
+#include "tensorflow/contrib/lite/kernels/internal/common.h"
#include "tensorflow/contrib/lite/kernels/activation_functor.h"
#include "tensorflow/contrib/lite/kernels/internal/common.h"
#include "tensorflow/contrib/lite/kernels/internal/optimized/tensor_utils_impl.h"
import zipfile
import numpy as np
from six import StringIO
+from six.moves import xrange
# TODO(aselle): Disable GPU for now
os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
with self.test_session() as sess:
init = init_ops.constant_initializer(0.5)
- with variable_scope.variable_scope("root", initializer=init):
+ with variable_scope.variable_scope("root",
+ initializer=init):
x = array_ops.zeros([1, 2])
c0 = array_ops.zeros([1, 2])
h0 = array_ops.zeros([1, 2])
xout, sout = cell()(x, state0)
sess.run([variables.global_variables_initializer()])
- res = sess.run(
- [xout, sout], {
- x.name: np.array([[1., 1.]]),
- c0.name: 0.1 * np.asarray([[0, 1]]),
- h0.name: 0.1 * np.asarray([[2, 3]]),
- })
+ res = sess.run([xout, sout], {
+ x.name: np.array([[1., 1.]]),
+ c0.name: 0.1 * np.asarray([[0, 1]]),
+ h0.name: 0.1 * np.asarray([[2, 3]]),
+ })
actual_state_c = res[1].c
actual_state_h = res[1].h
"""Tests cell w/o peepholes and w/o normalisation."""
def cell():
- return contrib_rnn_cell.WeightNormLSTMCell(
- 2, norm=False, use_peepholes=False)
+ return contrib_rnn_cell.WeightNormLSTMCell(2,
+ norm=False,
+ use_peepholes=False)
actual_c, actual_h = self._cell_output(cell)
"""Tests cell with peepholes and w/o normalisation."""
def cell():
- return contrib_rnn_cell.WeightNormLSTMCell(
- 2, norm=False, use_peepholes=True)
+ return contrib_rnn_cell.WeightNormLSTMCell(2,
+ norm=False,
+ use_peepholes=True)
actual_c, actual_h = self._cell_output(cell)
"""Tests cell w/o peepholes and with normalisation."""
def cell():
- return contrib_rnn_cell.WeightNormLSTMCell(
- 2, norm=True, use_peepholes=False)
+ return contrib_rnn_cell.WeightNormLSTMCell(2,
+ norm=True,
+ use_peepholes=False)
actual_c, actual_h = self._cell_output(cell)
"""Tests cell with peepholes and with normalisation."""
def cell():
- return contrib_rnn_cell.WeightNormLSTMCell(
- 2, norm=True, use_peepholes=True)
+ return contrib_rnn_cell.WeightNormLSTMCell(2,
+ norm=True,
+ use_peepholes=True)
actual_c, actual_h = self._cell_output(cell)
eos_token,
vocab_size,
dtype=probs.dtype,
- on_value=0.,
+ on_value=ops.convert_to_tensor(0., dtype=probs.dtype),
off_value=probs.dtype.min)
finished_probs = array_ops.tile(
array_ops.reshape(finished_row, [1, 1, -1]),
the data will be cycled through indefinitely.
num_readers: a integer, number of Readers to create.
reader_kwargs: an optional dict, of kwargs for the reader.
- shuffle: boolean, wether should shuffle the files and the records by using
+ shuffle: boolean, whether should shuffle the files and the records by using
RandomShuffleQueue as common_queue.
dtypes: A list of types. The length of dtypes must equal the number
of elements in each record. If it is None it will default to
void PackToProto(FertileSlot* slot) const override;
void InitLeafClassStats(int best_split_index, LeafStat* left_stats,
- LeafStat* right_stats) const;
+ LeafStat* right_stats) const override;
protected:
void ClassificationAddSplitStats() override {
tf_custom_op_library(
name = "python/ops/_trt_engine_op.so",
- srcs = ["ops/trt_engine_op.cc"],
+ srcs = [
+ "ops/trt_calib_op.cc",
+ "ops/trt_engine_op.cc",
+ ],
deps = [
":trt_engine_op_kernel",
":trt_shape_function",
cc_library(
name = "trt_engine_op_kernel",
- srcs = ["kernels/trt_engine_op.cc"],
- hdrs = ["kernels/trt_engine_op.h"],
+ srcs = [
+ "kernels/trt_calib_op.cc",
+ "kernels/trt_engine_op.cc",
+ ],
+ hdrs = [
+ "kernels/trt_calib_op.h",
+ "kernels/trt_engine_op.h",
+ ],
copts = tf_copts(),
deps = [
":trt_logging",
+ ":trt_resources",
"//tensorflow/core:gpu_headers_lib",
"//tensorflow/core:lib_proto_parsing",
"//tensorflow/core:stream_executor_headers_lib",
)
tf_gen_op_libs(
- op_lib_names = ["trt_engine_op"],
+ op_lib_names = [
+ "trt_engine_op",
+ "trt_calib_op",
+ ],
deps = if_tensorrt([
"@local_config_tensorrt//:nv_infer",
]),
name = "trt_engine_op",
gen_locally = True,
deps = [
+ ":trt_calib_op_op_lib",
":trt_engine_op_op_lib",
":trt_logging",
":trt_shape_function",
],
)
+tf_cuda_library(
+ name = "trt_resources",
+ srcs = [
+ "resources/trt_int8_calibrator.cc",
+ "resources/trt_resource_manager.cc",
+ ],
+ hdrs = [
+ "resources/trt_int8_calibrator.h",
+ "resources/trt_resource_manager.h",
+ "resources/trt_resources.h",
+ ],
+ deps = [
+ ":trt_logging",
+ "//tensorflow/core:framework_headers_lib",
+ "//tensorflow/core:framework_lite",
+ "//tensorflow/core:lib_proto_parsing",
+ ] + if_tensorrt([
+ "@local_config_tensorrt//:nv_infer",
+ ]),
+)
+
# Library for the node-level conversion portion of TensorRT operation creation
tf_cuda_library(
name = "trt_conversion",
deps = [
":segment",
":trt_logging",
+ ":trt_resources",
"//tensorflow/core/grappler:grappler_item",
"//tensorflow/core/grappler:utils",
"//tensorflow/core:framework",
CHECK_EQ_TYPE(tensor_r->getType(), dtype);
auto op_pair = ops.find(node_def.op());
if (op_pair == ops.end())
- return tensorflow::errors::Unimplemented(
- "binary op: " + node_def.op() +
- " not supported at: " + node_def.name());
+ return tensorflow::errors::Unimplemented("binary op: " + node_def.op() +
+ " not supported at: " +
+ node_def.name());
nvinfer1::IElementWiseLayer* layer = ctx.network()->addElementWise(
*const_cast<nvinfer1::ITensor*>(tensor_l),
<< std::to_string(op_info_vec.size());
// TODO(ben,jie): update TRT input format/dimension
- nvinfer1::DimsCHW input_dim_pseudo_chw;
- for (int i = 0; i < 3; i++) input_dim_pseudo_chw.d[i] = 1;
+ nvinfer1::DimsCHW input_dim_psuedo_chw;
+ for (int i = 0; i < 3; i++) input_dim_psuedo_chw.d[i] = 1;
for (int i = 1; i < op_info.shape().dim_size(); i++) {
VLOG(2) << "dimension: " << i
<< " , size: " << op_info.shape().dim(i).size();
- input_dim_pseudo_chw.d[i - 1] = op_info.shape().dim(i).size();
+ input_dim_psuedo_chw.d[i - 1] = op_info.shape().dim(i).size();
}
// TODO(ben,jie): proper way to restore input tensor name?
input_tensor_name = node_name + ":" + std::to_string(output_idx);
nvinfer1::ITensor* input_tensor = converter.network()->addInput(
- input_tensor_name.c_str(), dtype, input_dim_pseudo_chw);
+ input_tensor_name.c_str(), dtype, input_dim_psuedo_chw);
if (!input_tensor)
return tensorflow::errors::InvalidArgument(
--- /dev/null
+/* Copyright 2018 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/contrib/tensorrt/kernels/trt_calib_op.h"
+#include "tensorflow/contrib/tensorrt/resources/trt_int8_calibrator.h"
+#include "tensorflow/contrib/tensorrt/resources/trt_resource_manager.h"
+#include "tensorflow/contrib/tensorrt/resources/trt_resources.h"
+#include "tensorflow/core/framework/tensor.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/framework/tensor_types.h"
+#include "tensorflow/core/framework/types.h"
+
+#if GOOGLE_CUDA
+#if GOOGLE_TENSORRT
+#include "cuda_runtime_api.h"
+#include "tensorrt/include/NvInfer.h"
+
+namespace tensorflow {
+namespace tensorrt {
+
+TRTCalibOp::TRTCalibOp(OpKernelConstruction* context) : OpKernel(context) {
+ OP_REQUIRES_OK(context, context->GetAttr("segment_nodes", &segment_nodes_));
+ OP_REQUIRES_OK(context, context->GetAttr("input_names", &input_names_));
+ OP_REQUIRES_OK(context, context->GetAttr("resource_name", &resource_name_));
+};
+
+#define TYPECASE(dt, X, Y) \
+ case dt: { \
+ return (void*)X->flat<tensorflow::EnumToDataType<dt>::Type>().data(); \
+ }
+
+void* GetTensorAddress(const Tensor* tensor_ptr) {
+ auto tensor_type = tensor_ptr->dtype();
+ switch (tensor_type) {
+ TYPECASE(tensorflow::DT_FLOAT, tensor_ptr, dest_ptr);
+ TYPECASE(tensorflow::DT_HALF, tensor_ptr, dest_ptr);
+ TYPECASE(tensorflow::DT_INT8, tensor_ptr, dest_ptr);
+ default: {
+ LOG(FATAL) << "Unsupported Data type "
+ << tensorflow::DataTypeString(tensor_type);
+ return nullptr;
+ }
+ }
+}
+
+void TRTCalibOp::Compute(tensorflow::OpKernelContext* ctx) {
+ // TODO(aaroey): make sure ctx->resource_mgr() is used in future PR.
+ auto trt_rm = tensorflow::tensorrt::TRTResourceManager::instance();
+ auto res_mgr = trt_rm->getManager("TRTCalibOps");
+ tensorflow::tensorrt::TRTCalibrationResource* calib_res = nullptr;
+ auto status = res_mgr->Lookup(resource_name_, resource_name_, &calib_res);
+
+ if (!status.ok()) {
+ ctx->SetStatus(status);
+ return;
+ }
+ int num_inputs = ctx->num_inputs();
+ // first run instantiate calibrator
+ if (calib_res->calibrator_ == nullptr) {
+ dev_tensors_.resize(num_inputs);
+ int batch_size = ctx->input(0).dim_size(0);
+ VLOG(1) << " Constructing calibrator";
+ for (int i = 0; i < num_inputs; i++) {
+ // allocate workspace on device for inputs
+ const tensorflow::Tensor& t = ctx->input(i);
+ OP_REQUIRES_OK(ctx,
+ ctx->allocate_persistent(t.dtype(), t.shape(),
+ &dev_tensors_.at(i), nullptr));
+ const auto device_tensor = dev_tensors_.at(i).AccessTensor(ctx);
+ CHECK_EQ(t.TotalBytes(), device_tensor->TotalBytes());
+ void* device_address = GetTensorAddress(device_tensor);
+ device_buffers_.emplace(input_names_.at(i),
+ std::pair<void*, size_t>(
+ device_address, device_tensor->TotalBytes()));
+ }
+
+ calib_res->calibrator_ =
+ new TRTInt8Calibrator(device_buffers_, batch_size, resource_name_);
+ string label(resource_name_);
+ calib_res->thr_ = new std::thread([calib_res, label]() {
+ VLOG(1) << "Starting calibration thread, Calibration Resource @ "
+ << calib_res;
+ calib_res->builder_->setInt8Calibrator(calib_res->calibrator_);
+ calib_res->builder_->setInt8Mode(true);
+ calib_res->engine_ = calib_res->builder_->buildCudaEngine(
+ *calib_res->network_); // will loop until we terminate calibrator
+ VLOG(1) << "Calibration loop terminated " << label;
+ });
+ VLOG(1) << "initialized calibrator resource";
+ } // calibrator initialized
+
+ // Pass input data to calibrator
+ std::unordered_map<string, void*> input_data;
+ for (int i = 0; i < num_inputs; i++) {
+ const Tensor& t = ctx->input(i);
+ void* data_address = GetTensorAddress(&t);
+ const auto device_tensor = dev_tensors_.at(i).AccessTensor(ctx);
+ CHECK_EQ(t.TotalBytes(),
+ device_tensor->TotalBytes()); // use the tensor so FW keeps it
+ input_data.emplace(input_names_.at(i), data_address);
+ ctx->set_output(i, t);
+ }
+ VLOG(2) << "Filled map for sending";
+ calib_res->calibrator_->setBatch(input_data);
+ VLOG(2) << "Passed calibration data";
+ // TODO(aaroey): make sure we wait for the completion of calibration on the
+ // last batch in future PR.
+};
+
+#undef TYPECASE
+
+REGISTER_KERNEL_BUILDER(Name("TRTCalibOp").Device(DEVICE_GPU), TRTCalibOp);
+
+} // namespace tensorrt
+} // namespace tensorflow
+#endif
+#endif
--- /dev/null
+/* Copyright 2018 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.
+==============================================================================*/
+
+#ifndef TENSORFLOW_CONTRIB_TENSORRT_KERNELS_TRT_CALIB_OP_H
+#define TENSORFLOW_CONTRIB_TENSORRT_KERNELS_TRT_CALIB_OP_H
+
+#include <memory>
+#include <string>
+#include <unordered_map>
+#include <utility>
+#include <vector>
+#include "tensorflow/core/framework/op.h"
+#include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/framework/tensor_shape.h"
+#include "tensorflow/core/platform/types.h"
+
+#if GOOGLE_CUDA
+#if GOOGLE_TENSORRT
+namespace tensorflow {
+namespace tensorrt {
+// TODO(sami): Convert this to async kernel!
+class TRTCalibOp : public OpKernel {
+ public:
+ explicit TRTCalibOp(OpKernelConstruction* context);
+
+ void Compute(OpKernelContext* context) override;
+
+ private:
+ string resource_name_;
+ std::vector<string> segment_nodes_;
+ std::vector<string> input_names_;
+ std::vector<tensorflow::TensorShape> shapes_;
+ std::unordered_map<string, std::pair<void*, size_t>> device_buffers_;
+ std::vector<tensorflow::PersistentTensor> dev_tensors_;
+};
+} // namespace tensorrt
+} // namespace tensorflow
+#endif
+#endif
+#endif // TENSORFLOW_CONTRIB_TENSORRT_KERNELS_TRT_CALIB_OP_H
--- /dev/null
+/* Copyright 2018 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.h"
+#include "tensorflow/core/framework/shape_inference.h"
+namespace tensorflow {
+
+REGISTER_OP("TRTCalibOp")
+ .Attr("segment_nodes: list(string)") // names of the ops in segment
+ .Attr("segment_output_names: list(string)") // names of the output ops in
+ // segment
+ .Attr("input_names: list(string)") // names of the inputs for
+ // passing into tensorrt
+ .Attr("resource_name: string")
+ .Attr("InT: list({int8, float16, float32})")
+ .Input("in_tensor: InT")
+ .Output("out_tensor: InT")
+ .SetShapeFn([](tensorflow::shape_inference::InferenceContext* c) {
+ for (int i = 0; i < c->num_inputs(); i++) {
+ c->set_output(i, c->input(i));
+ }
+ return Status::OK();
+ });
+
+} // namespace tensorflow
--- /dev/null
+/* Copyright 2018 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/contrib/tensorrt/resources/trt_int8_calibrator.h"
+
+#include <atomic>
+#include <chrono>
+#include <unordered_map>
+
+#include "tensorflow/core/platform/logging.h"
+
+#if GOOGLE_CUDA
+#if GOOGLE_TENSORRT
+#include "cuda_runtime_api.h"
+
+namespace tensorflow {
+namespace tensorrt {
+
+// set the batch size before constructing the thread to execute engine
+int TRTInt8Calibrator::getBatchSize() const { return batch_size_; }
+
+TRTInt8Calibrator::TRTInt8Calibrator(
+ const std::unordered_map<string, std::pair<void*, size_t>>& dev_buffers,
+ int batch_size, string engine_name)
+ : batch_size_(batch_size),
+ done_(false),
+ dev_buffers_(dev_buffers),
+ calib_running_(false),
+ engine_name_(engine_name) {}
+
+bool TRTInt8Calibrator::setBatch(
+ const std::unordered_map<string, void*>& data) {
+ // TODO(aaroey): make sure that in future PR:
+ // 1. the mutex_lock is outside of the loop
+ // 2. wait() is used instead of wait_for()
+ // 3. done_ is to be protected by the mutex
+ // 4. the first batch is not missed
+ if (done_) return false;
+ while (calib_running_.load(
+ std::memory_order_acquire)) { // wait while calibration is running
+ tensorflow::mutex_lock l(cond_mtx_);
+ cond_.wait_for(l, std::chrono::milliseconds(50));
+ if (done_) return false;
+ }
+ VLOG(1) << "Set Batch Waiting finished";
+ for (const auto it : data) {
+ auto devptr = dev_buffers_.find(it.first);
+ if (devptr == dev_buffers_.end()) {
+ LOG(FATAL) << "FATAL " << engine_name_ << " input name '" << it.first
+ << "' does not match with the buffer names";
+ }
+ const auto& d = devptr->second;
+
+ // TODO(aaroey): we should not use sync copy on default stream. Make sure
+ // stream->ThenMemcpy() is used in future PRs.
+ auto status =
+ cudaMemcpy(d.first, it.second, d.second, cudaMemcpyDeviceToDevice);
+ if (status != cudaSuccess) {
+ LOG(FATAL) << "cudaMemcpy " << engine_name_ << " for '" << it.first
+ << "' failed with " << status;
+ }
+ }
+ calib_running_.store(true, std::memory_order_release); // release builder
+ cond_.notify_all();
+ return true;
+}
+
+bool TRTInt8Calibrator::getBatch(void** bindings, const char** names,
+ int num_bindings) {
+ calib_running_.store(false, std::memory_order_release); // wait for new batch
+ cond_.notify_all();
+ while (!calib_running_.load(
+ std::memory_order_acquire)) { // wait until new batch arrives
+ tensorflow::mutex_lock l(cond_mtx_);
+ cond_.wait_for(l, std::chrono::milliseconds(50));
+ if (done_) return false;
+ }
+ if (done_) {
+ return false;
+ }
+
+ for (int i = 0; i < num_bindings; i++) {
+ auto it = dev_buffers_.find(names[i]);
+ if (it == dev_buffers_.end()) {
+ LOG(FATAL) << "Calibration engine asked for unknown tensor name '"
+ << names[i] << "' at position " << i;
+ }
+
+ bindings[i] = it->second.first;
+ }
+ return true;
+}
+
+const void* TRTInt8Calibrator::readCalibrationCache(std::size_t& length) {
+ return nullptr;
+}
+
+void TRTInt8Calibrator::writeCalibrationCache(const void* ptr,
+ std::size_t length) {}
+TRTInt8Calibrator::~TRTInt8Calibrator() {
+ VLOG(1) << "Destroying calibrator for " << engine_name_;
+}
+
+} // namespace tensorrt
+} // namespace tensorflow
+#endif
+#endif
--- /dev/null
+/* Copyright 2018 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.
+==============================================================================*/
+
+#ifndef TENSORFLOW_CONTRIB_TENSORRT_RESOURCES_TRT_INT8_CALIBRATOR_H_
+#define TENSORFLOW_CONTRIB_TENSORRT_RESOURCES_TRT_INT8_CALIBRATOR_H_
+
+#include <atomic>
+#include <string>
+#include <unordered_map>
+#include <utility>
+#include "tensorflow/core/platform/mutex.h"
+
+#if GOOGLE_CUDA
+#if GOOGLE_TENSORRT
+#include "tensorrt/include/NvInfer.h"
+namespace tensorflow {
+namespace tensorrt {
+// This class provides a 1 element queue to match TFs push model to
+// TRTs pull model for calibration. When TRT implements a means for
+// a push calibration This class should be updated accordingly
+
+struct TRTInt8Calibrator : public nvinfer1::IInt8EntropyCalibrator {
+ public:
+ TRTInt8Calibrator(
+ const std::unordered_map<string, std::pair<void*, size_t>>& dev_buffers,
+ int batch_size, string engine_name);
+ int getBatchSize() const override;
+ bool getBatch(void* bindings[], const char* names[],
+ int num_bindings) override;
+ bool setBatch(const std::unordered_map<string, void*>& data);
+ void setDone() { done_ = true; }
+ const void* readCalibrationCache(std::size_t& length) override;
+ void writeCalibrationCache(const void* ptr, std::size_t length) override;
+ ~TRTInt8Calibrator();
+
+ private:
+ const int batch_size_;
+ tensorflow::mutex cond_mtx_; // mutex for condition_variable
+ tensorflow::condition_variable cond_; // condition variable to implement
+ // producer-consumer queue for
+ // calibration
+ bool done_;
+ const std::unordered_map<string, std::pair<void*, size_t>>
+ dev_buffers_; // map to keep tensorrt input buffers and sizes keyed with
+ // buffer names
+ std::atomic_bool calib_running_;
+ string engine_name_;
+};
+} // namespace tensorrt
+} // namespace tensorflow
+#endif // TENSORFLOW_CONTRIB_TENSORRT_RESOURCES_TRT_INT8_CALIBRATOR_H_
+#endif
+#endif
--- /dev/null
+/* Copyright 2018 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/contrib/tensorrt/resources/trt_resource_manager.h"
+#include "tensorflow/core/platform/logging.h"
+
+namespace tensorflow {
+namespace tensorrt {
+
+std::shared_ptr<tensorflow::ResourceMgr>
+tensorflow::tensorrt::TRTResourceManager::getManager(const string& op_name) {
+ // mutex is held for lookup only. Most instantiations where mutex will be held
+ // longer will be during op creation and should be ok.
+ tensorflow::mutex_lock lock(map_mutex_);
+ auto s = managers_.find(op_name);
+ if (s == managers_.end()) {
+ auto it = managers_.emplace(
+ op_name, std::make_shared<tensorflow::ResourceMgr>(op_name));
+ VLOG(1) << "Returning a new manager " << op_name;
+ return it.first->second;
+ }
+ VLOG(1) << "Returning old manager " << op_name;
+ return s->second;
+}
+
+} // namespace tensorrt
+} // namespace tensorflow
--- /dev/null
+/* Copyright 2018 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.
+==============================================================================*/
+
+#ifndef TENSORFLOW_CONTRIB_TENSORRT_RESOURCES_TRT_RESOURCE_MANAGER_H_
+#define TENSORFLOW_CONTRIB_TENSORRT_RESOURCES_TRT_RESOURCE_MANAGER_H_
+#include <memory>
+
+#include <string>
+#include <unordered_map>
+#include "tensorflow/core/framework/resource_mgr.h"
+#include "tensorflow/core/platform/mutex.h"
+
+namespace tensorflow {
+namespace tensorrt {
+
+class TRTResourceManager {
+ TRTResourceManager() = default;
+
+ public:
+ static std::shared_ptr<TRTResourceManager> instance() {
+ static std::shared_ptr<TRTResourceManager> instance_(
+ new TRTResourceManager);
+ return instance_;
+ }
+ // returns a manager for given op, if it doesn't exists it creates one
+ std::shared_ptr<tensorflow::ResourceMgr> getManager(const string& op_name);
+
+ private:
+ std::unordered_map<string, std::shared_ptr<tensorflow::ResourceMgr>>
+ managers_;
+ tensorflow::mutex map_mutex_;
+};
+
+} // namespace tensorrt
+} // namespace tensorflow
+
+#endif // TENSORFLOW_CONTRIB_TENSORRT_RESOURCE_TRT_RESOURCE_MANAGER_H_
--- /dev/null
+/* Copyright 2018 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.
+==============================================================================*/
+
+#ifndef TENSORFLOW_CONTRIB_TENSORRT_RESOURCES_TRTRESOURCES_H_
+#define TENSORFLOW_CONTRIB_TENSORRT_RESOURCES_TRTRESOURCES_H_
+
+#include <list>
+#include <sstream>
+#include <string>
+#include <thread>
+#include <vector>
+#include "tensorflow/contrib/tensorrt/log/trt_logger.h"
+#include "tensorflow/core/framework/resource_mgr.h"
+
+#if GOOGLE_CUDA
+#if GOOGLE_TENSORRT
+#include "tensorflow/contrib/tensorrt/resources/trt_int8_calibrator.h"
+#include "tensorrt/include/NvInfer.h"
+
+namespace tensorflow {
+namespace tensorrt {
+class TRTCalibrationResource : public tensorflow::ResourceBase {
+ public:
+ TRTCalibrationResource()
+ : calibrator_(nullptr),
+ builder_(nullptr),
+ network_(nullptr),
+ engine_(nullptr),
+ logger_(nullptr),
+ thr_(nullptr) {}
+ string DebugString() override {
+ std::stringstream oss;
+ oss << " Calibrator = " << std::hex << calibrator_ << std::dec << std::endl
+ << " Builder = " << std::hex << builder_ << std::dec << std::endl
+ << " Network = " << std::hex << network_ << std::dec << std::endl
+ << " Engine = " << std::hex << engine_ << std::dec << std::endl
+ << " Logger = " << std::hex << logger_ << std::dec << std::endl
+ << " Thread = " << std::hex << thr_ << std::dec << std::endl;
+ return oss.str();
+ }
+ ~TRTCalibrationResource() {
+ VLOG(0) << "Destroying Calibration Resource " << std::endl << DebugString();
+ }
+ TRTInt8Calibrator* calibrator_;
+ nvinfer1::IBuilder* builder_;
+ nvinfer1::INetworkDefinition* network_;
+ nvinfer1::ICudaEngine* engine_;
+ tensorflow::tensorrt::Logger* logger_;
+ // TODO(sami): Use threadpool threads!
+ std::thread* thr_;
+};
+
+class TRTWeightStore : public tensorflow::ResourceBase {
+ public:
+ TRTWeightStore() {}
+ std::list<std::vector<uint8_t>> store_;
+ string DebugString() override {
+ std::stringstream oss;
+ size_t lenBytes = 0;
+ for (const auto& v : store_) {
+ lenBytes += v.size() * sizeof(uint8_t);
+ }
+ oss << " Number of entries = " << store_.size() << std::endl
+ << " Total number of bytes = "
+ << store_.size() * sizeof(std::vector<uint8_t>) + lenBytes << std::endl;
+ return oss.str();
+ }
+ virtual ~TRTWeightStore() { VLOG(1) << "Destroying store" << DebugString(); }
+};
+
+class TRTEngineResource : public tensorflow::ResourceBase {
+ public:
+ TRTEngineResource() : runtime_(nullptr), ctx_(nullptr){};
+ string DebugString() override { return string(""); }
+ nvinfer1::IRuntime* runtime_;
+ nvinfer1::IExecutionContext* ctx_;
+};
+
+} // namespace tensorrt
+} // namespace tensorflow
+#endif // TENSORFLOW_CONTRIB_TENSORRT_RESOURCEMGR_TRTRESOURCES_H_
+#endif
+#endif
"//tensorflow/python:framework_ops",
"//tensorflow/python:math_ops",
"//tensorflow/python:state_ops",
+ "//tensorflow/python:summary",
"//tensorflow/python:util",
"//tensorflow/python:variable_scope",
"//tensorflow/python/estimator:estimator_py",
"//tensorflow/python/estimator:export",
"//tensorflow/python/estimator:head",
+ "//tensorflow/python/estimator:metric_keys",
],
)
from tensorflow.python.estimator import estimator_lib
from tensorflow.python.estimator.canned import head as head_lib
+from tensorflow.python.estimator.canned import metric_keys
from tensorflow.python.estimator.export import export_lib
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.ops import state_ops
from tensorflow.python.ops import variable_scope
from tensorflow.python.util import nest
+from tensorflow.python.summary import summary
def time_series_regression_head(model,
self.input_statistics_generator = input_statistics_generator
self._name = name
+ @property
+ def name(self):
+ return self._name
+
+ # TODO(terrytangyuan): consolidate `model_outputs` and `_Head.LossSpec`
+ # once `_Head.create_loss` becomes extendable
+ def create_loss(self, features, mode, logits=None, labels=None):
+ """See `_Head`."""
+ model_outputs = self.state_manager.define_loss(
+ self.model, features, mode)
+ summary.scalar(
+ head_lib._summary_key(self._name, metric_keys.MetricKeys.LOSS),
+ model_outputs.loss)
+ return model_outputs
+
+ @property
+ def logits_dimension(self):
+ """See `_Head`."""
+ return 1
+
def _train_ops(self, features):
"""Add training ops to the graph."""
+ mode = estimator_lib.ModeKeys.TRAIN
with variable_scope.variable_scope(
"model",
# Use ResourceVariables to avoid race conditions.
use_resource=True):
- model_outputs = self.state_manager.define_loss(
- self.model, features, estimator_lib.ModeKeys.TRAIN)
+ model_outputs = self.create_loss(features, mode)
train_op = optimizers.optimize_loss(
model_outputs.loss,
learning_rate=None)
return estimator_lib.EstimatorSpec(
loss=model_outputs.loss,
- mode=estimator_lib.ModeKeys.TRAIN,
+ mode=mode,
train_op=train_op)
- # TODO(terrytangyuan): suffix summary and metrics keys by `"/" + name`
- @property
- def name(self):
- return self._name
-
- # TODO(terrytangyuan): unused for now. Need to decouple
- # `state_manager.define_loss` to satisfy the extendable return signature of
- # `_Head.create_loss`.
- def create_loss(self, features, mode, logits, labels):
- """See `_Head`."""
- return None
-
- # TODO(terrytangyuan): check label dimension
- @property
- def logits_dimension(self):
- return None
-
def _evaluate_ops(self, features):
"""Add ops for evaluation (aka filtering) to the graph."""
+ mode = estimator_lib.ModeKeys.EVAL
with variable_scope.variable_scope("model", use_resource=True):
- model_outputs = self.state_manager.define_loss(
- self.model, features, estimator_lib.ModeKeys.EVAL)
+ model_outputs = self.create_loss(features, mode)
metrics = {}
# Just output in-sample predictions for the last chunk seen
for prediction_key, prediction_value in model_outputs.predictions.items():
model_outputs.end_state))
return estimator_lib.EstimatorSpec(
loss=model_outputs.loss,
- mode=estimator_lib.ModeKeys.EVAL,
+ mode=mode,
eval_metric_ops=metrics,
predictions={})
with variable_scope.variable_scope("model", use_resource=True):
prediction_outputs = self.model.predict(features=features)
with variable_scope.variable_scope("model", reuse=True):
- filtering_outputs = self.state_manager.define_loss(
- self.model, features, estimator_lib.ModeKeys.EVAL)
-
+ filtering_outputs = self.create_loss(
+ features, estimator_lib.ModeKeys.EVAL)
return estimator_lib.EstimatorSpec(
mode=estimator_lib.ModeKeys.PREDICT,
export_outputs={
def create_estimator_spec(self, features, mode, labels=None):
"""Performs basic error checking and returns an EstimatorSpec."""
- with ops.name_scope("head"):
+ with ops.name_scope(self._name, "head"):
if labels:
raise ValueError(
"The model received a `labels` dictionary, which is "
1. When the sender receives a tensor request, the source tensor may or may not be ready yet. The situation is handled through a process of tag matching:
* If the request arrives before the tensor is ready, then a callback is put in a local table, and will be invoked once the tensor arrives.
- * If the tensor is ready before the request arives, than the tensor is put in a local table. When the request arrives, it will invoke the callback immediately.
+ * If the tensor is ready before the request arrives, than the tensor is put in a local table. When the request arrives, it will invoke the callback immediately.
In code it is done by calling **RecvLocalAsync()**, which receives the tensor's key, step-id, and the callback.
2. When the callback is invoked, the relevant tensor is removed from the tag matching table. In the case where we need to send the tensor's meta-data, the **RdmaTensorResponse** will store a copy of the tensor until the re-request arrives.
3. The sending of protocol messages (**RDMA_MESSAGE_TENSOR_REQUEST**, **RDMA_MESSAGE_META_DATA_RESPONSE** and **RDMA_MESSAGE_TENSOR_RE_REQUEST**) is done by the class **RdmaMessageBuffer**. All messages are sent using RDMA writes from/to fixed messages buffers. This implies that we cannot send on a specific channel more than one message at a time. In order to synchronize the messages, the **RdmaMessageBuffer** holds the a local and remote buffer statuses which can be either busy or idle. When a write is issued, both statuses will be changed to busy. When the write-complete event is received, the local status is changed to idle. When the write is received on the remote side, the remote side will parse the message, and return an ACK back to the sending side on which the sending side will update the remote status to idle. When both the local and remote statuses are idle, the next message can be sent.
* type - The message type.
* request_index - Request index.
* is_dead/data_type/tensor_shape/tensor_bytes - The up-to-date meta-data.
-* **RDMA_MESSAGE_BUFFER_RESPONSE** - (receiver ==> sender) Tensor re-requset after meta-data update and reallocation of result/proxy tensors.
+* **RDMA_MESSAGE_BUFFER_RESPONSE** - (receiver ==> sender) Tensor re-request after meta-data update and reallocation of result/proxy tensors.
* type - The message type.
* name (name_size) - Name of the requested tensor.
* step_id - Step ID.
#include "tensorflow/core/distributed_runtime/rendezvous_mgr_interface.h"
#include "tensorflow/core/distributed_runtime/rpc/grpc_util.h"
#include "tensorflow/core/distributed_runtime/session_mgr.h"
+#include "tensorflow/core/distributed_runtime/rpc/grpc_util.h"
#include "tensorflow/core/framework/rendezvous.h"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/lib/core/status.h"
--- /dev/null
+op {
+ graph_op_name: "UniqueWithCountsV2"
+ in_arg {
+ name: "x"
+ description: <<END
+A `Tensor`.
+END
+ }
+ in_arg {
+ name: "axis"
+ description: <<END
+A `Tensor` of type `int32` (default: None). The axis of the Tensor to
+find the unique elements.
+END
+ }
+ out_arg {
+ name: "y"
+ description: <<END
+A `Tensor`. Unique elements along the `axis` of `Tensor` x.
+END
+ }
+ out_arg {
+ name: "idx"
+ description: <<END
+A 1-D Tensor. Has the same type as x that contains the index of each
+value of x in the output y.
+END
+ }
+ out_arg {
+ name: "count"
+ description: <<END
+A 1-D Tensor. The count of each value of x in the output y.
+END
+ }
+ summary: "Finds unique elements along an axis of a tensor."
+ description: <<END
+This operation either returns a tensor `y` containing unique elements
+along the `axis` of a tensor. The returned unique elements is sorted
+in the same order as they occur along `axis` in `x`.
+This operation also returns a tensor `idx` and a tensor `count`
+that are the same size as the number of the elements in `x` along the
+`axis` dimension. The `idx` contains the index in the unique output `y`
+and the `count` contains the count in the unique output `y`.
+In other words, for an `1-D` tensor `x` with `axis = None:
+
+`y[idx[i]] = x[i] for i in [0, 1,...,rank(x) - 1]`
+
+For example:
+
+```
+# tensor 'x' is [1, 1, 2, 4, 4, 4, 7, 8, 8]
+y, idx, count = unique_with_counts(x)
+y ==> [1, 2, 4, 7, 8]
+idx ==> [0, 0, 1, 2, 2, 2, 3, 4, 4]
+count ==> [2, 1, 3, 1, 2]
+```
+
+For an `2-D` tensor `x` with `axis = 0`:
+
+```
+# tensor 'x' is [[1, 0, 0],
+# [1, 0, 0],
+# [2, 0, 0]]
+y, idx, count = unique_with_counts(x, axis=0)
+y ==> [[1, 0, 0],
+ [2, 0, 0]]
+idx ==> [0, 0, 1]
+count ==> [2, 1]
+```
+
+For an `2-D` tensor `x` with `axis = 1`:
+
+```
+# tensor 'x' is [[1, 0, 0],
+# [1, 0, 0],
+# [2, 0, 0]]
+y, idx, count = unique_with_counts(x, axis=1)
+y ==> [[1, 0],
+ [1, 0],
+ [2, 0]]
+idx ==> [0, 1, 1]
+count ==> [1, 2]
+```
+END
+}
has size `num_segments`.
END
}
- summary: "Computes the Max along segments of a tensor."
+ summary: "Computes the maximum along segments of a tensor."
description: <<END
Read @{$math_ops#Segmentation$the section on segmentation} for an explanation of
segments.
-This operator is similar to the [unsorted segment sum operator](../../../api_docs/python/math_ops.md#UnsortedSegmentSum).
-Instead of computing the sum over segments, it computes the maximum
-such that:
+This operator is similar to the unsorted segment sum operator found
+[(here)](../../../api_docs/python/math_ops.md#UnsortedSegmentSum).
+Instead of computing the sum over segments, it computes the maximum such that:
\\(output_i = \max_j data_j\\) where max is over `j` such
that `segment_ids[j] == i`.
-If the maximum is empty for a given segment ID `i`, it outputs the smallest possible value for specific numeric type,
- `output[i] = numeric_limits<T>::min()`.
+If the maximum is empty for a given segment ID `i`, it outputs the smallest
+possible value for the specific numeric type,
+`output[i] = numeric_limits<T>::lowest()`.
<div style="width:70%; margin:auto; margin-bottom:10px; margin-top:20px;">
<img style="width:100%" src="https://www.tensorflow.org/images/UnsortedSegmentMax.png" alt>
--- /dev/null
+op {
+ graph_op_name: "UnsortedSegmentMin"
+ in_arg {
+ name: "segment_ids"
+ description: <<END
+A 1-D tensor whose rank is equal to the rank of `data`'s
+first dimension.
+END
+ }
+ out_arg {
+ name: "output"
+ description: <<END
+Has same shape as data, except for dimension 0 which
+has size `num_segments`.
+END
+ }
+ summary: "Computes the minimum along segments of a tensor."
+ description: <<END
+Read @{$math_ops#segmentation$the section on segmentation} for an explanation of
+segments.
+
+This operator is similar to the unsorted segment sum operator found
+[(here)](../../../api_docs/python/math_ops.md#UnsortedSegmentSum).
+Instead of computing the sum over segments, it computes the minimum such that:
+
+\\(output_i = \min_j data_j\\) where min is over `j` such
+that `segment_ids[j] == i`.
+
+If the minimum is empty for a given segment ID `i`, it outputs the largest
+possible value for the specific numeric type,
+`output[i] = numeric_limits<T>::max()`.
+END
+}
--- /dev/null
+op {
+ graph_op_name: "UnsortedSegmentProd"
+ in_arg {
+ name: "segment_ids"
+ description: <<END
+A 1-D tensor whose rank is equal to the rank of `data`'s
+first dimension.
+END
+ }
+ out_arg {
+ name: "output"
+ description: <<END
+Has same shape as data, except for dimension 0 which
+has size `num_segments`.
+END
+ }
+ summary: "Computes the product along segments of a tensor."
+ description: <<END
+Read @{$math_ops#segmentation$the section on segmentation} for an explanation of
+segments.
+
+This operator is similar to the unsorted segment sum operator found
+[(here)](../../../api_docs/python/math_ops.md#UnsortedSegmentSum).
+Instead of computing the sum over segments, it computes the product of all
+entries belonging to a segment such that:
+
+\\(output_i = \prod_j data_j\\) where the product is over `j` such
+that `segment_ids[j] == i`.
+
+If there is no entry for a given segment ID `i`, it outputs 1.
+END
+}
--- /dev/null
+op {
+ graph_op_name: "UniqueWithCounts"
+ visibility: HIDDEN
+}
--- /dev/null
+op {
+ graph_op_name: "UniqueWithCountsV2"
+ visibility: HIDDEN
+}
const TensorReferenceVector& tensor_refs) override;
Status FillContextMap(const Graph* graph,
- DeviceContextMap* device_context_map);
+ DeviceContextMap* device_context_map) override;
void Compute(OpKernel* op_kernel, OpKernelContext* context) override;
worker_cache_factory_(std::move(worker_cache_factory)) {}
string SessionMgr::WorkerNameFromServerDef(const ServerDef& server_def) {
- return strings::StrCat("/job:", server_def.job_name(),
- "/replica:0/task:", server_def.task_index());
+ return strings::StrCat("/job:", server_def.job_name(), "/replica:0/task:",
+ server_def.task_index());
}
Status SessionMgr::CreateSession(const string& session,
#define TENSORFLOW_FRAMEWORK_NUMERIC_TYPES_H_
#include <complex>
-
#include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor"
// Disable clang-format to prevent 'FixedPoint' header from being included
// before 'Tensor' header on which it depends.
} // namespace tensorflow
+
+
+
+static inline tensorflow::bfloat16 FloatToBFloat16(float float_val) {
+#if __BYTE_ORDER__ == __ORDER_BIG_ENDIAN__
+ return *reinterpret_cast<tensorflow::bfloat16*>(
+ reinterpret_cast<uint16_t*>(&float_val));
+#else
+ return *reinterpret_cast<tensorflow::bfloat16*>(
+ &(reinterpret_cast<uint16_t*>(&float_val)[1]));
+#endif
+}
+
namespace Eigen {
-// TOOD(xpan): We probably need to overwrite more methods to have correct eigen
-// behavior. E.g. loest(), is_integer, etc. See NumTraits.h in eigen.
+// TODO(xpan): We probably need to overwrite more methods to have correct eigen
+// behavior. E.g. epsilon(), dummy_precision, etc. See NumTraits.h in eigen.
template <>
struct NumTraits<tensorflow::bfloat16>
- : GenericNumTraits<tensorflow::bfloat16> {};
+ : GenericNumTraits<tensorflow::bfloat16> {
+ enum {
+ IsInteger = 0,
+ IsSigned = 1,
+ RequireInitialization = 0
+ };
+ static EIGEN_STRONG_INLINE tensorflow::bfloat16 highest() {
+ return FloatToBFloat16(NumTraits<float>::highest());
+ }
+
+ static EIGEN_STRONG_INLINE tensorflow::bfloat16 lowest() {
+ return FloatToBFloat16(NumTraits<float>::lowest());
+ }
+
+ static EIGEN_STRONG_INLINE tensorflow::bfloat16 infinity() {
+ return FloatToBFloat16(NumTraits<float>::infinity());
+ }
+
+ static EIGEN_STRONG_INLINE tensorflow::bfloat16 quiet_NaN() {
+ return FloatToBFloat16(NumTraits<float>::quiet_NaN());
+ }
+};
+
using ::tensorflow::operator==;
using ::tensorflow::operator!=;
return errors::Internal(
"No unary variant binary_op function found for binary variant op "
"enum: ",
- op, " Variant type_name: '", a.TypeName(),
- "' for device type: ", device);
+ op, " Variant type_name: '", a.TypeName(), "' for device type: ",
+ device);
}
return (*binary_op_fn)(ctx, a, b, out);
}
"//tensorflow/core/grappler:op_types",
"//tensorflow/core/grappler:utils",
"//tensorflow/core/grappler/costs:graph_properties",
+ "//tensorflow/core/grappler/utils:frame",
],
)
name = "loop_optimizer_test",
size = "small",
srcs = ["loop_optimizer_test.cc"],
+ tags = [
+ "manual",
+ "no_oss", # b/74111495
+ "notap",
+ ],
deps = [
":loop_optimizer",
"//tensorflow/cc:cc_ops",
#include "tensorflow/core/grappler/optimizers/loop_optimizer.h"
+#include <algorithm>
+#include <limits>
#include <unordered_map>
#include <unordered_set>
+#include <vector>
+#include <deque>
#include "tensorflow/core/framework/attr_value.pb.h"
+#include "tensorflow/core/framework/node_def.pb.h"
#include "tensorflow/core/framework/op.h"
+#include "tensorflow/core/framework/types.h"
#include "tensorflow/core/grappler/costs/graph_properties.h"
#include "tensorflow/core/grappler/grappler_item.h"
#include "tensorflow/core/grappler/op_types.h"
#include "tensorflow/core/grappler/optimizers/constant_folding.h"
#include "tensorflow/core/grappler/utils.h"
+#include "tensorflow/core/grappler/utils/frame.h"
#include "tensorflow/core/lib/core/errors.h"
#include "tensorflow/core/lib/core/stringpiece.h"
#include "tensorflow/core/lib/strings/strcat.h"
+#include "tensorflow/core/platform/tensor_coding.h"
+#include "tensorflow/core/util/device_name_utils.h"
+#include "tensorflow/core/util/saved_tensor_slice_util.h"
+
+using tensorflow::strings::StrCat;
namespace tensorflow {
namespace grappler {
} // namespace
+Status LoopOptimizer::LINMHandleInvariantEnter(NodeDef* node,
+ const int num_outputs) {
+ auto consumers = node_map_->GetOutputs(node->name());
+ std::vector<string> enter_control_inputs;
+ string enter_input;
+ for (auto& input : node->input()) {
+ if (IsControlInput(input)) {
+ enter_control_inputs.push_back(input);
+ } else {
+ enter_input = input;
+ }
+ }
+ for (auto* consumer : consumers) {
+ if (invariant_nodes_.count(consumer)) {
+ for (int i = 0; i < consumer->input_size(); ++i) {
+ if (NodeName(consumer->input(i)) == node->name()) {
+ consumer->set_input(i, enter_input);
+ node_map_->AddOutput(NodeName(enter_input), consumer->name());
+ node_map_->RemoveOutput(node->name(), consumer->name());
+ }
+ }
+ for (auto& control_input : enter_control_inputs) {
+ consumer->add_input(control_input);
+ node_map_->AddOutput(NodeName(control_input), consumer->name());
+ }
+ }
+ }
+ return Status::OK();
+}
+
+Status LoopOptimizer::LINMHandleConst(NodeDef* node,
+ const int num_outputs, const int frame_id) {
+ NodeDef* const_node;
+ if (num_outputs == 0) {
+ // all successor nodes are invariant
+ // Remove the control inputs from this frame to the const node,
+ // when moving it out of the frame (in parent frame)
+ const_node = node;
+ node_map_->RemoveInputs(node->name());
+ node->clear_input();
+ } else {
+ // some successor nodes are variant
+ // Have to keep the const node in the frame,
+ // so create a new one outside the frame (in parent frame)
+ const_node = optimized_graph_->add_node();
+ const_node->set_name(AddPrefixToNodeName(node->name(), kLoopOptimizer));
+ const_node->set_op("Const");
+ const_node->set_device(node->device());
+ *const_node->mutable_attr() = node->attr();
+ node_map_->AddNode(const_node->name(), const_node);
+ auto consumers = node_map_->GetOutputs(node->name());
+ for (auto* consumer : consumers) {
+ if (invariant_nodes_.count(consumer)) {
+ for (int i = 0; i < consumer->input_size(); ++i) {
+ if (NodeName(consumer->input(i)) == node->name()) {
+ if (IsControlInput(consumer->input(i))) {
+ *consumer->mutable_input(i) = AsControlDependency(*const_node);
+ } else {
+ *consumer->mutable_input(i) = const_node->name();
+ }
+ node_map_->AddOutput(const_node->name(), consumer->name());
+ node_map_->RemoveOutput(node->name(), consumer->name());
+ }
+ }
+ }
+ }
+ }
+ // add a control input from the parent frame
+ auto parent_it = frame_parent_.find(frame_id);
+ if (parent_it != frame_parent_.end()) {
+ int parent_id = parent_it->second;
+ auto loop_cond_it = loop_cond_.find(parent_id);
+ if (loop_cond_it == loop_cond_.end()) {
+ return errors::InvalidArgument(
+ "Frame ", frame_id, " doesn't have a LoopCond node");
+ }
+ auto& loop_cond_name = loop_cond_it->second->name();
+ NodeDef* switch_node = nullptr;
+ for (auto* node : node_map_->GetOutputs(loop_cond_name)) {
+ if (node->op() == "Switch") {
+ switch_node = node;
+ break;
+ }
+ }
+ if (!switch_node) {
+ return errors::InvalidArgument(
+ "LoopCond node of Frame ", frame_id,
+ " doesn't connect to any Switch node");
+ }
+ string switch_output = StrCat(switch_node->name(), ":1");
+ const string ctrl_dep = ConstantFolding::AddControlDependency(
+ switch_output, optimized_graph_, node_map_.get());
+ const_node->add_input(ctrl_dep);
+ node_map_->AddOutput(NodeName(ctrl_dep), const_node->name());
+ }
+ return Status::OK();
+}
+
+Status LoopOptimizer::LINMHandleInvariantNode(NodeDef* node,
+ const int num_outputs, const int frame_id) {
+ // have to remove control inputs to the invariant node from the same frame
+ // when moving this node out of this frame
+ for (int i = 0; i < node->input_size(); ++i) {
+ if (IsControlInput(node->input(i))) {
+ node->mutable_input()->SwapElements(i, node->input_size() - 1);
+ node->mutable_input()->RemoveLast();
+ }
+ }
+ if (num_outputs == 0) {
+ return Status::OK();
+ }
+
+ DataTypeVector input_types;
+ DataTypeVector output_types;
+ OpRegistryInterface* op_registry = OpRegistry::Global();
+ const OpRegistrationData* op_reg_data = nullptr;
+ TF_RETURN_IF_ERROR(
+ op_registry->LookUp(node->op(), &op_reg_data));
+ TF_RETURN_IF_ERROR(
+ InOutTypesForNode(*node, op_reg_data->op_def,
+ &input_types, &output_types));
+
+ auto consumers = node_map_->GetOutputs(node->name());
+ string fname = invariant_enters_[frame_id][0]->attr().at("frame_name").s();
+ int piterations = invariant_enters_[frame_id][0]
+ ->attr().at("parallel_iterations").i();
+ for (auto* consumer : consumers) {
+ if (!invariant_nodes_.count(consumer)) {
+ for (int i = 0; i < consumer->input_size(); ++i) {
+ int port;
+ string node_name = ParseNodeName(consumer->input(i), &port);
+ if (node_name != node->name()) {
+ continue;
+ }
+ if (port < 0) {
+ return errors::InvalidArgument(
+ "Invariant node should not have control outputs "
+ "to variant node");
+ }
+ DataType output_type = output_types[port];
+ NodeDef* new_enter = optimized_graph_->add_node();
+ new_enter->set_op("Enter");
+ new_enter->set_device(node->device());
+ new_enter->set_name(AddPrefixToNodeName(
+ StrCat(fname, "_enter_", new_enter_id_++), kLoopOptimizer));
+ AttrValue data_type;
+ data_type.set_type(output_type);
+ new_enter->mutable_attr()->insert({"T", data_type});
+ AttrValue frame_name;
+ frame_name.set_s(fname);
+ new_enter->mutable_attr()->insert({"frame_name", frame_name});
+ AttrValue is_const;
+ is_const.set_b(true);
+ new_enter->mutable_attr()->insert({"is_constant", is_const});
+ AttrValue parallel_iterations;
+ parallel_iterations.set_i(piterations);
+ new_enter->mutable_attr()->insert(
+ {"parallel_iterations", parallel_iterations});
+ new_enter->add_input(consumer->input(i));
+ *consumer->mutable_input(i) = new_enter->name();
+ node_map_->AddNode(new_enter->name(), new_enter);
+ node_map_->AddOutput(node->name(), new_enter->name());
+ node_map_->AddOutput(new_enter->name(), consumer->name());
+ }
+ }
+ }
+ return Status::OK();
+}
+
+Status LoopOptimizer::MoveInvariantNodes(const int frame_id) {
+ for (auto iter = invariant_nodes_.begin();
+ iter != invariant_nodes_.end(); ++iter) {
+ auto* invariant_node = iter->first;
+ const int num_outputs = iter->second;
+ if (IsEnter(*invariant_node)) {
+ TF_RETURN_IF_ERROR(
+ LINMHandleInvariantEnter(invariant_node, num_outputs));
+ } else if (IsConstant(*invariant_node)) {
+ TF_RETURN_IF_ERROR(
+ LINMHandleConst(invariant_node, num_outputs, frame_id));
+ } else {
+ TF_RETURN_IF_ERROR(
+ LINMHandleInvariantNode(invariant_node, num_outputs, frame_id));
+ }
+ }
+ return Status::OK();
+}
+
+Status LoopOptimizer::RevertInvariantNodes() {
+ std::deque<const NodeDef*> reverted_nodes;
+ for (auto iter=invariant_nodes_.begin(); iter != invariant_nodes_.end();) {
+ bool erased = false;
+ const auto* node = iter->first;
+ if (!IsConstant(*node) && !IsEnter(*node) && iter->second > 0) {
+ auto& consumers = node_map_->GetOutputs(node->name());
+ for (auto* consumer : consumers) {
+ if (!invariant_nodes_.count(consumer)) {
+ for (const auto& input : consumer->input()) {
+ if (IsControlInput(input) && NodeName(input) == node->name()) {
+ reverted_nodes.push_back(node);
+ invariant_nodes_.erase(iter++);
+ erased = true;
+ break;
+ }
+ }
+ if (erased) break;
+ }
+ }
+ }
+ if (!erased) ++iter;
+ }
+ while (!reverted_nodes.empty()) {
+ const auto* node = reverted_nodes.front();
+ reverted_nodes.pop_front();
+ std::set<NodeDef*> producers;
+ for (const auto& input : node->input()) {
+ auto* producer = node_map_->GetNode(input);
+ auto iter = invariant_nodes_.find(producer);
+ if (iter != invariant_nodes_.end()) {
+ if (IsControlInput(input) &&
+ !IsConstant(*producer) && !IsEnter(*producer)) {
+ reverted_nodes.push_back(producer);
+ invariant_nodes_.erase(iter);
+ } else {
+ producers.insert(producer);
+ }
+ }
+ }
+ for (auto* producer : producers) {
+ auto iter = invariant_nodes_.find(producer);
+ if (iter != invariant_nodes_.end()) {
+ ++iter->second;
+ }
+ }
+ for (auto* consumer : node_map_->GetOutputs(node->name())) {
+ auto iter = invariant_nodes_.find(consumer);
+ if (iter != invariant_nodes_.end()) {
+ reverted_nodes.push_back(consumer);
+ invariant_nodes_.erase(iter);
+ }
+ }
+ }
+ return Status::OK();
+}
+
+Status LoopOptimizer::FindInvariantNodes(NodeDef* node) {
+ auto consumers = node_map_->GetOutputs(node->name());
+ invariant_nodes_.insert(std::make_pair(node, consumers.size()));
+ for (auto* consumer : consumers) {
+ if (invariant_nodes_.count(consumer) ||
+ ModifiesFrameInfo(*consumer)) {
+ continue;
+ }
+ bool is_invariant = true;
+ for (const auto& input : consumer->input()) {
+ if (!IsControlInput(input)) {
+ const auto& name = NodeName(input);
+ auto* producer = node_map_->GetNode(name);
+ if (!invariant_nodes_.count(producer)) {
+ if (IsConstant(*producer)) {
+ invariant_nodes_.insert(
+ std::make_pair(producer, node_map_->GetOutputs(name).size()));
+ } else {
+ is_invariant = false;
+ break;
+ }
+ }
+ }
+ }
+ if (is_invariant) {
+ std::set<NodeDef*> producers;
+ for (const auto& input : consumer->input()) {
+ auto* producer = node_map_->GetNode(input);
+ producers.insert(producer);
+ }
+ for (auto* producer : producers) {
+ auto iter = invariant_nodes_.find(producer);
+ if (iter != invariant_nodes_.end()) {
+ --iter->second;
+ }
+ }
+ TF_RETURN_IF_ERROR(FindInvariantNodes(consumer));
+ }
+ }
+ return Status::OK();
+}
+
+Status LoopOptimizer::LoopInvariantNodeMotion() {
+ std::deque<int> worklist;
+ for (auto iter = frame_map_.begin(); iter != frame_map_.end(); ++iter) {
+ auto* node = iter->first;
+ auto& frame_ids = iter->second;
+ if (frame_ids.size() >= 3) {
+ for (unsigned int i = 1; i < frame_ids.size() - 1; ++i) {
+ frame_parent_[frame_ids[i]] = frame_ids[i - 1];
+ frame_children_[frame_ids[i]].insert(frame_ids[i + 1]);
+ }
+ }
+ if (frame_ids.size() >= 2) {
+ frame_children_[frame_ids[0]].insert(frame_ids[1]);
+ frame_parent_[frame_ids.back()] = frame_ids[frame_ids.size() - 2];
+ }
+ if (!frame_ids.empty()) {
+ frame_children_.insert(std::make_pair(frame_ids.back(), empty_set_));
+ if (node->op() == "LoopCond") {
+ if (loop_cond_.count(frame_ids.back())) {
+ return errors::InvalidArgument(
+ "Loop ", frame_ids.back(),
+ " has more than one LoopCond node: ", node->name(), " and ",
+ loop_cond_[frame_ids.back()]->name());
+ }
+ loop_cond_[frame_ids.back()] = node;
+ }
+ if (IsEnter(*node) && node->attr().at("is_constant").b()) {
+ invariant_enters_[frame_ids.back()].push_back(
+ const_cast<NodeDef*>(node));
+ }
+ }
+ }
+
+ for (auto it = frame_children_.begin(); it != frame_children_.end(); ++it) {
+ if (it->second.empty()) {
+ worklist.push_back(it->first);
+ }
+ }
+
+ while (!worklist.empty()) {
+ int frame_id = worklist.front();
+ new_enter_id_ = 0;
+ worklist.pop_front();
+ auto parent_it = frame_parent_.find(frame_id);
+ if (parent_it != frame_parent_.end()) {
+ int parent_id = parent_it->second;
+ frame_children_[parent_id].erase(frame_id);
+ if (frame_children_[parent_id].empty()) {
+ worklist.push_back(parent_id);
+ }
+ }
+
+ if (invariant_enters_[frame_id].empty()) {
+ continue;
+ }
+ invariant_nodes_.clear();
+ for (auto* enter : invariant_enters_[frame_id]) {
+ TF_RETURN_IF_ERROR(FindInvariantNodes(enter));
+ }
+
+ // revert invariant nodes that have control outputs to variant nodes
+ TF_RETURN_IF_ERROR(RevertInvariantNodes());
+
+ TF_RETURN_IF_ERROR(MoveInvariantNodes(frame_id));
+ }
+ return Status::OK();
+}
+
Status LoopOptimizer::Optimize(Cluster* cluster, const GrapplerItem& item,
GraphDef* optimized_graph) {
- Status status = RemoveStackOps(item.graph, optimized_graph);
- return status;
+ TF_RETURN_IF_ERROR(RemoveStackOps(item.graph, optimized_graph));
+
+ optimized_graph_ = optimized_graph;
+
+ // Set up helper data structures.
+ node_map_.reset(new NodeMap(optimized_graph_));
+ int num_frames;
+ TF_RETURN_IF_ERROR(IdentifyFramesWithNodeMap(*optimized_graph_, *node_map_,
+ &frame_map_, &num_frames));
+
+ TF_RETURN_IF_ERROR(LoopInvariantNodeMotion());
+ return Status::OK();
}
void LoopOptimizer::Feedback(Cluster* /*cluster*/, const GrapplerItem& /*item*/,
#define TENSORFLOW_CORE_GRAPPLER_OPTIMIZERS_LOOP_OPTIMIZER_H_
#include <unordered_set>
+#include "tensorflow/core/grappler/costs/graph_properties.h"
#include "tensorflow/core/grappler/optimizers/graph_optimizer.h"
#include "tensorflow/core/grappler/utils.h"
+#include "tensorflow/core/grappler/utils/frame.h"
#include "tensorflow/core/protobuf/rewriter_config.pb.h"
namespace tensorflow {
namespace grappler {
+constexpr char kLoopOptimizer[] = "LoopOptimizer";
+
class LoopOptimizer : public GraphOptimizer {
public:
LoopOptimizer() : opt_level_(RewriterConfig::ON) {}
const GraphDef& optimized_graph, double result) override;
private:
+ Status LoopInvariantNodeMotion();
+ Status FindInvariantNodes(NodeDef* node);
+ Status RevertInvariantNodes();
+ Status MoveInvariantNodes(const int frame_id);
+ Status LINMHandleInvariantNode(NodeDef* node, const int num_outputs,
+ const int frame_id);
+ Status LINMHandleConst(NodeDef* node, const int num_outputs,
+ const int frame_id);
+ Status LINMHandleInvariantEnter(NodeDef* node, const int num_outputs);
+
+ std::map<NodeDef*, int> invariant_nodes_;
+ std::set<int> empty_set_;
+ std::map<int, std::set<int>> frame_children_;
+ std::map<int, int> frame_parent_;
+ std::map<int, const NodeDef*> loop_cond_;
+ std::map<int, std::vector<NodeDef*>> invariant_enters_;
+ int new_enter_id_;
RewriterConfig::Toggle opt_level_;
+
+ std::unique_ptr<NodeMap> node_map_;
+ FrameMap frame_map_;
+ std::unique_ptr<GraphProperties> graph_properties_;
+ GraphDef* optimized_graph_; // Not owned.
};
} // end namespace grappler
namespace grappler {
namespace {
-class LoopOptimizerTest : public ::testing::Test {};
+class LoopOptimizerTest : public ::testing::Test {
+ protected:
+ static NodeDef CreateNode(const string& name,
+ const std::vector<string>& inputs) {
+ return CreateNode(name, "Identity", "", false, 0, inputs);
+ }
+ static NodeDef CreateNode(const string& name, const string& op,
+ const std::vector<string>& inputs) {
+ return CreateNode(name, op, "", false, 0, inputs);
+ }
+ static NodeDef CreateNode(const string& name, const string& op,
+ const string& frame,
+ const bool is_constant,
+ const int piterations,
+ const std::vector<string>& inputs) {
+ NodeDef node;
+ node.set_name(name);
+ if (!op.empty()) {
+ node.set_op(op);
+ }
+ if (!frame.empty()) {
+ AttrValue frame_name;
+ frame_name.set_s(frame);
+ node.mutable_attr()->insert({"frame_name", frame_name});
+ }
+ if (op == "Enter") {
+ AttrValue is_const;
+ is_const.set_b(is_constant);
+ node.mutable_attr()->insert({"is_constant", is_const});
+ AttrValue parallel_iterations;
+ parallel_iterations.set_i(piterations);
+ node.mutable_attr()->insert(
+ {"parallel_iterations", parallel_iterations});
+ }
+ AttrValue type;
+ type.set_type(DT_FLOAT);
+ node.mutable_attr()->insert({"T", type});
+ for (const string& input : inputs) {
+ node.add_input(input);
+ }
+ return node;
+ }
+};
+
+TEST_F(LoopOptimizerTest, Basic) {
+ GraphDef graph;
+ *graph.add_node() = CreateNode("0", {});
+ *graph.add_node() = CreateNode(
+ "InvariantEnter", "Enter", "while/while_context", true, 1, {"0"});
+ *graph.add_node() = CreateNode(
+ "InvariantAdd", "Add", {"InvariantEnter", "InvariantEnter"});
+ *graph.add_node() = CreateNode(
+ "VariantAdd", "Add", {"InvariantAdd", "Identity"});
+ *graph.add_node() = CreateNode(
+ "VariantEnter", "Enter", "while/while_context", false, 1, {"0"});
+ *graph.add_node() = CreateNode(
+ "Merge", "Merge", {"VariantEnter", "NextIteration"});
+ *graph.add_node() = CreateNode("Less/y", "Const", {"^Identity"});
+ *graph.add_node() = CreateNode("Less", "Less", {"VariantAdd", "less/y"});
+ *graph.add_node() = CreateNode("LoopCond", "LoopCond", {"Less"});
+ *graph.add_node() = CreateNode("Switch", "Switch", {"Merge", "LoopCond"});
+ *graph.add_node() = CreateNode("Identity", {"Switch:1"});
+ *graph.add_node() = CreateNode(
+ "NextIteration", "NextIteration", {"VariantAdd"});
+ *graph.add_node() = CreateNode("Exit", "Exit", {"Switch"});
+ *graph.add_node() = CreateNode("1", {"Exit"});
+
+ GrapplerItem item;
+ item.graph = graph;
+
+ LoopOptimizer optimizer;
+ GraphDef output;
+ Status status = optimizer.Optimize(nullptr, item, &output);
+ TF_EXPECT_OK(status);
+
+ std::unique_ptr<NodeMap> node_map;
+ std::unordered_map<const NodeDef*, std::vector<int>> frames;
+ int num_frames;
+
+ node_map.reset(new NodeMap(&graph));
+ EXPECT_TRUE(IdentifyFrames(graph, &frames, &num_frames).ok());
+ EXPECT_EQ(num_frames, 1);
+ EXPECT_EQ(frames.at(node_map->GetNode("InvariantAdd")).size(), 1);
+ EXPECT_EQ(frames.at(node_map->GetNode("InvariantAdd")).back(), 0);
+ EXPECT_EQ(frames.at(node_map->GetNode("VariantAdd")).size(), 1);
+ EXPECT_EQ(frames.at(node_map->GetNode("VariantAdd")).back(), 0);
+
+ node_map.reset(new NodeMap(&output));
+ EXPECT_TRUE(IdentifyFrames(output, &frames, &num_frames).ok());
+ EXPECT_EQ(num_frames, 1);
+ EXPECT_EQ(frames.at(node_map->GetNode("InvariantAdd")).size(), 0);
+ EXPECT_EQ(frames.at(node_map->GetNode("VariantAdd")).size(), 1);
+ EXPECT_EQ(frames.at(node_map->GetNode("VariantAdd")).back(), 0);
+}
+
+TEST_F(LoopOptimizerTest, Const) {
+ GraphDef graph;
+ *graph.add_node() = CreateNode("0", {});
+ *graph.add_node() = CreateNode(
+ "InvariantEnter", "Enter", "while/while_context", true, 1, {"0"});
+ *graph.add_node() = CreateNode("Const", "Const", {"^Identity"});
+ *graph.add_node() = CreateNode(
+ "InvariantAdd", "Add", {"InvariantEnter", "Const"});
+ *graph.add_node() = CreateNode(
+ "VariantAdd", "Add", {"InvariantAdd", "Identity"});
+ *graph.add_node() = CreateNode(
+ "VariantEnter", "Enter", "while/while_context", false, 1, {"0"});
+ *graph.add_node() = CreateNode(
+ "Merge", "Merge", {"VariantEnter", "NextIteration"});
+ *graph.add_node() = CreateNode("Less/y", "Const", {"^Identity"});
+ *graph.add_node() = CreateNode("Less", "Less", {"VariantAdd", "less/y"});
+ *graph.add_node() = CreateNode("LoopCond", "LoopCond", {"Less"});
+ *graph.add_node() = CreateNode("Switch", "Switch", {"Merge", "LoopCond"});
+ *graph.add_node() = CreateNode("Identity", {"Switch:1"});
+ *graph.add_node() = CreateNode(
+ "NextIteration", "NextIteration", {"VariantAdd"});
+ *graph.add_node() = CreateNode("Exit", "Exit", {"Switch"});
+ *graph.add_node() = CreateNode("1", {"Exit"});
+
+ GrapplerItem item;
+ item.graph = graph;
+
+ LoopOptimizer optimizer;
+ GraphDef output;
+ Status status = optimizer.Optimize(nullptr, item, &output);
+ TF_EXPECT_OK(status);
+
+ std::unique_ptr<NodeMap> node_map;
+ std::unordered_map<const NodeDef*, std::vector<int>> frames;
+ int num_frames;
+
+ node_map.reset(new NodeMap(&graph));
+ EXPECT_TRUE(IdentifyFrames(graph, &frames, &num_frames).ok());
+ EXPECT_EQ(num_frames, 1);
+ EXPECT_EQ(frames.at(node_map->GetNode("InvariantAdd")).size(), 1);
+ EXPECT_EQ(frames.at(node_map->GetNode("InvariantAdd")).back(), 0);
+ EXPECT_EQ(frames.at(node_map->GetNode("Const")).size(), 1);
+ EXPECT_EQ(frames.at(node_map->GetNode("Const")).back(), 0);
+
+ node_map.reset(new NodeMap(&output));
+ EXPECT_TRUE(IdentifyFrames(output, &frames, &num_frames).ok());
+ EXPECT_EQ(num_frames, 1);
+ EXPECT_EQ(frames.at(node_map->GetNode("InvariantAdd")).size(), 0);
+ EXPECT_EQ(frames.at(node_map->GetNode("Const")).size(), 0);
+}
+
+TEST_F(LoopOptimizerTest, ControlOutput) {
+ GraphDef graph;
+ *graph.add_node() = CreateNode("0", {});
+ *graph.add_node() = CreateNode(
+ "InvariantEnter", "Enter", "while/while_context", true, 1, {"0"});
+ *graph.add_node() = CreateNode(
+ "InvariantAdd", "Add", {"InvariantEnter", "InvariantEnter"});
+ *graph.add_node() = CreateNode(
+ "VariantAdd", "Add", {"InvariantAdd", "Identity"});
+ *graph.add_node() = CreateNode(
+ "VariantEnter", "Enter", "while/while_context", false, 1, {"0"});
+ *graph.add_node() = CreateNode(
+ "Merge", "Merge", {"VariantEnter", "NextIteration"});
+ *graph.add_node() = CreateNode("Less/y", "Const", {"^Identity"});
+ *graph.add_node() = CreateNode(
+ "Less", "Less", {"VariantAdd", "less/y", "^InvariantAdd"});
+ *graph.add_node() = CreateNode("LoopCond", "LoopCond", {"Less"});
+ *graph.add_node() = CreateNode("Switch", "Switch", {"Merge", "LoopCond"});
+ *graph.add_node() = CreateNode("Identity", {"Switch:1"});
+ *graph.add_node() = CreateNode(
+ "NextIteration", "NextIteration", {"VariantAdd"});
+ *graph.add_node() = CreateNode("Exit", "Exit", {"Switch"});
+ *graph.add_node() = CreateNode("1", {"Exit"});
+
+ GrapplerItem item;
+ item.graph = graph;
+
+ LoopOptimizer optimizer;
+ GraphDef output;
+ Status status = optimizer.Optimize(nullptr, item, &output);
+ TF_EXPECT_OK(status);
+
+ std::unique_ptr<NodeMap> node_map;
+ std::unordered_map<const NodeDef*, std::vector<int>> frames;
+ int num_frames;
+
+ node_map.reset(new NodeMap(&graph));
+ EXPECT_TRUE(IdentifyFrames(graph, &frames, &num_frames).ok());
+ EXPECT_EQ(num_frames, 1);
+ EXPECT_EQ(frames.at(node_map->GetNode("InvariantAdd")).size(), 1);
+ EXPECT_EQ(frames.at(node_map->GetNode("InvariantAdd")).back(), 0);
+
+ node_map.reset(new NodeMap(&output));
+ EXPECT_TRUE(IdentifyFrames(output, &frames, &num_frames).ok());
+ EXPECT_EQ(num_frames, 1);
+ EXPECT_EQ(frames.at(node_map->GetNode("InvariantAdd")).size(), 1);
+ EXPECT_EQ(frames.at(node_map->GetNode("InvariantAdd")).back(), 0);
+}
+
+TEST_F(LoopOptimizerTest, NestedLoop1) {
+ GraphDef graph;
+ *graph.add_node() = CreateNode("0", {});
+ *graph.add_node() = CreateNode(
+ "InvariantEnter", "Enter", "while/while_context", true, 1, {"0"});
+ *graph.add_node() = CreateNode(
+ "InvariantAdd", "Add", {"InvariantEnter", "InvariantEnter"});
+ *graph.add_node() = CreateNode(
+ "VariantAdd", "Add", {"InvariantAdd", "Identity"});
+ *graph.add_node() = CreateNode(
+ "VariantEnter", "Enter", "while/while_context", false, 1, {"0"});
+ *graph.add_node() = CreateNode(
+ "Merge", "Merge", {"VariantEnter", "NextIteration"});
+ *graph.add_node() = CreateNode("Less/y", "Const", {"^Identity"});
+ *graph.add_node() = CreateNode("Less", "Less", {"Exit2", "less/y"});
+ *graph.add_node() = CreateNode("LoopCond", "LoopCond", {"Less"});
+ *graph.add_node() = CreateNode("Switch", "Switch", {"Merge", "LoopCond"});
+ *graph.add_node() = CreateNode("Identity", {"Switch:1"});
+ *graph.add_node() = CreateNode(
+ "NextIteration", "NextIteration", {"Exit2"});
+ *graph.add_node() = CreateNode("Exit", "Exit", {"Switch"});
+ *graph.add_node() = CreateNode("1", {"Exit"});
+
+ *graph.add_node() = CreateNode(
+ "InvariantEnter2", "Enter", "while/while/while_context", true, 1,
+ {"VariantAdd"});
+ *graph.add_node() = CreateNode(
+ "InvariantAdd2", "Add", {"InvariantEnter2", "InvariantEnter2"});
+ *graph.add_node() = CreateNode(
+ "VariantAdd2", "Add", {"InvariantAdd2", "Identity2"});
+ *graph.add_node() = CreateNode(
+ "VariantEnter2", "Enter", "while/while/while_context", false, 1,
+ {"VariantEnter"});
+ *graph.add_node() = CreateNode(
+ "Merge2", "Merge", {"VariantEnter2", "NextIteration2"});
+ *graph.add_node() = CreateNode("Less2/y", "Const", {"^Identity2"});
+ *graph.add_node() = CreateNode("Less2", "Less", {"VariantAdd2", "less2/y"});
+ *graph.add_node() = CreateNode("LoopCond2", "LoopCond", {"Less2"});
+ *graph.add_node() = CreateNode("Switch2", "Switch", {"Merge2", "LoopCond2"});
+ *graph.add_node() = CreateNode("Identity2", {"Switch2:1"});
+ *graph.add_node() = CreateNode(
+ "NextIteration2", "NextIteration", {"VariantAdd2"});
+ *graph.add_node() = CreateNode("Exit2", "Exit", {"Switch2"});
+
+ GrapplerItem item;
+ item.graph = graph;
+
+ LoopOptimizer optimizer;
+ GraphDef output;
+ Status status = optimizer.Optimize(nullptr, item, &output);
+ TF_EXPECT_OK(status);
+
+ std::unique_ptr<NodeMap> node_map;
+ std::unordered_map<const NodeDef*, std::vector<int>> frames;
+ int num_frames;
+
+ node_map.reset(new NodeMap(&graph));
+ EXPECT_TRUE(IdentifyFrames(graph, &frames, &num_frames).ok());
+ EXPECT_EQ(num_frames, 2);
+ EXPECT_EQ(frames.at(node_map->GetNode("InvariantAdd2")).size(), 2);
+ EXPECT_EQ(frames.at(node_map->GetNode("InvariantAdd2")).back(), 1);
+ EXPECT_EQ(frames.at(node_map->GetNode("VariantAdd2")).size(), 2);
+ EXPECT_EQ(frames.at(node_map->GetNode("VariantAdd2")).back(), 1);
+ EXPECT_EQ(frames.at(node_map->GetNode("InvariantAdd")).size(), 1);
+ EXPECT_EQ(frames.at(node_map->GetNode("InvariantAdd")).back(), 0);
+
+ node_map.reset(new NodeMap(&output));
+ EXPECT_TRUE(IdentifyFrames(output, &frames, &num_frames).ok());
+ EXPECT_EQ(num_frames, 2);
+ EXPECT_EQ(frames.at(node_map->GetNode("InvariantAdd2")).size(), 1);
+ EXPECT_EQ(frames.at(node_map->GetNode("InvariantAdd2")).back(), 0);
+ EXPECT_EQ(frames.at(node_map->GetNode("VariantAdd2")).size(), 2);
+ EXPECT_EQ(frames.at(node_map->GetNode("VariantAdd2")).back(), 1);
+ EXPECT_EQ(frames.at(node_map->GetNode("InvariantAdd")).size(), 0);
+}
+
+TEST_F(LoopOptimizerTest, NestedLoop2) {
+ GraphDef graph;
+ *graph.add_node() = CreateNode("0", {});
+ *graph.add_node() = CreateNode(
+ "InvariantEnter", "Enter", "while/while_context", true, 1, {"0"});
+ *graph.add_node() = CreateNode(
+ "InvariantAdd", "Add", {"InvariantEnter", "InvariantEnter"});
+ *graph.add_node() = CreateNode(
+ "VariantAdd", "Add", {"InvariantAdd", "Identity"});
+ *graph.add_node() = CreateNode(
+ "VariantEnter", "Enter", "while/while_context", false, 1, {"0"});
+ *graph.add_node() = CreateNode(
+ "Merge", "Merge", {"VariantEnter", "NextIteration"});
+ *graph.add_node() = CreateNode("Less/y", "Const", {"^Identity"});
+ *graph.add_node() = CreateNode("Less", "Less", {"Exit2", "less/y"});
+ *graph.add_node() = CreateNode("LoopCond", "LoopCond", {"Less"});
+ *graph.add_node() = CreateNode("Switch", "Switch", {"Merge", "LoopCond"});
+ *graph.add_node() = CreateNode("Identity", {"Switch:1"});
+ *graph.add_node() = CreateNode(
+ "NextIteration", "NextIteration", {"Exit2"});
+ *graph.add_node() = CreateNode("Exit", "Exit", {"Switch"});
+ *graph.add_node() = CreateNode("1", {"Exit"});
+
+ *graph.add_node() = CreateNode(
+ "InvariantEnter2", "Enter", "while/while/while_context", true, 1,
+ {"InvariantAdd"});
+ *graph.add_node() = CreateNode(
+ "InvariantAdd2", "Add", {"InvariantEnter2", "InvariantEnter2"});
+ *graph.add_node() = CreateNode(
+ "VariantAdd2", "Add", {"InvariantAdd2", "Identity2"});
+ *graph.add_node() = CreateNode(
+ "VariantEnter2", "Enter", "while/while/while_context", false, 1,
+ {"VariantEnter"});
+ *graph.add_node() = CreateNode(
+ "Merge2", "Merge", {"VariantEnter2", "NextIteration2"});
+ *graph.add_node() = CreateNode("Less2/y", "Const", {"^Identity2"});
+ *graph.add_node() = CreateNode("Less2", "Less", {"VariantAdd2", "less2/y"});
+ *graph.add_node() = CreateNode("LoopCond2", "LoopCond", {"Less2"});
+ *graph.add_node() = CreateNode("Switch2", "Switch", {"Merge2", "LoopCond2"});
+ *graph.add_node() = CreateNode("Identity2", {"Switch2:1"});
+ *graph.add_node() = CreateNode(
+ "NextIteration2", "NextIteration", {"VariantAdd2"});
+ *graph.add_node() = CreateNode("Exit2", "Exit", {"Switch2"});
+
+ GrapplerItem item;
+ item.graph = graph;
+
+ LoopOptimizer optimizer;
+ GraphDef output;
+ Status status = optimizer.Optimize(nullptr, item, &output);
+ TF_EXPECT_OK(status);
+
+ std::unique_ptr<NodeMap> node_map;
+ std::unordered_map<const NodeDef*, std::vector<int>> frames;
+ int num_frames;
+
+ node_map.reset(new NodeMap(&graph));
+ EXPECT_TRUE(IdentifyFrames(graph, &frames, &num_frames).ok());
+ EXPECT_EQ(num_frames, 2);
+ EXPECT_EQ(frames.at(node_map->GetNode("InvariantAdd2")).size(), 2);
+ EXPECT_EQ(frames.at(node_map->GetNode("InvariantAdd2")).back(), 1);
+ EXPECT_EQ(frames.at(node_map->GetNode("VariantAdd2")).size(), 2);
+ EXPECT_EQ(frames.at(node_map->GetNode("VariantAdd2")).back(), 1);
+
+ node_map.reset(new NodeMap(&output));
+ EXPECT_TRUE(IdentifyFrames(output, &frames, &num_frames).ok());
+ EXPECT_EQ(num_frames, 2);
+ EXPECT_EQ(frames.at(node_map->GetNode("InvariantAdd2")).size(), 0);
+ EXPECT_EQ(frames.at(node_map->GetNode("VariantAdd2")).size(), 2);
+ EXPECT_EQ(frames.at(node_map->GetNode("VariantAdd2")).back(), 1);
+}
+
+TEST_F(LoopOptimizerTest, NestedLoopConst1) {
+ GraphDef graph;
+ *graph.add_node() = CreateNode("0", {});
+ *graph.add_node() = CreateNode(
+ "InvariantEnter", "Enter", "while/while_context", true, 1, {"0"});
+ *graph.add_node() = CreateNode(
+ "InvariantAdd", "Add", {"InvariantEnter", "InvariantEnter"});
+ *graph.add_node() = CreateNode(
+ "VariantAdd", "Add", {"InvariantAdd", "Identity"});
+ *graph.add_node() = CreateNode(
+ "VariantEnter", "Enter", "while/while_context", false, 1, {"0"});
+ *graph.add_node() = CreateNode(
+ "Merge", "Merge", {"VariantEnter", "NextIteration"});
+ *graph.add_node() = CreateNode("Less/y", "Const", {"^Identity"});
+ *graph.add_node() = CreateNode("Less", "Less", {"Exit2", "less/y"});
+ *graph.add_node() = CreateNode("LoopCond", "LoopCond", {"Less"});
+ *graph.add_node() = CreateNode("Switch", "Switch", {"Merge", "LoopCond"});
+ *graph.add_node() = CreateNode("Identity", {"Switch:1"});
+ *graph.add_node() = CreateNode(
+ "NextIteration", "NextIteration", {"Exit2"});
+ *graph.add_node() = CreateNode("Exit", "Exit", {"Switch"});
+ *graph.add_node() = CreateNode("1", {"Exit"});
+
+ *graph.add_node() = CreateNode(
+ "InvariantEnter2", "Enter", "while/while/while_context", true, 1,
+ {"VariantAdd"});
+ *graph.add_node() = CreateNode("Const2", "Const", {"^Identity2"});
+ *graph.add_node() = CreateNode(
+ "InvariantAdd2", "Add", {"InvariantEnter2", "Const2"});
+ *graph.add_node() = CreateNode(
+ "VariantAdd2", "Add", {"InvariantAdd2", "Identity2"});
+ *graph.add_node() = CreateNode(
+ "VariantEnter2", "Enter", "while/while/while_context", false, 1,
+ {"VariantEnter"});
+ *graph.add_node() = CreateNode(
+ "Merge2", "Merge", {"VariantEnter2", "NextIteration2"});
+ *graph.add_node() = CreateNode("Less2/y", "Const", {"^Identity2"});
+ *graph.add_node() = CreateNode("Less2", "Less", {"VariantAdd2", "less2/y"});
+ *graph.add_node() = CreateNode("LoopCond2", "LoopCond", {"Less2"});
+ *graph.add_node() = CreateNode("Switch2", "Switch", {"Merge2", "LoopCond2"});
+ *graph.add_node() = CreateNode("Identity2", {"Switch2:1"});
+ *graph.add_node() = CreateNode(
+ "NextIteration2", "NextIteration", {"VariantAdd2"});
+ *graph.add_node() = CreateNode("Exit2", "Exit", {"Switch2"});
+
+ GrapplerItem item;
+ item.graph = graph;
+
+ LoopOptimizer optimizer;
+ GraphDef output;
+ Status status = optimizer.Optimize(nullptr, item, &output);
+ TF_EXPECT_OK(status);
+
+ std::unique_ptr<NodeMap> node_map;
+ std::unordered_map<const NodeDef*, std::vector<int>> frames;
+ int num_frames;
+
+ node_map.reset(new NodeMap(&graph));
+ EXPECT_TRUE(IdentifyFrames(graph, &frames, &num_frames).ok());
+ EXPECT_EQ(num_frames, 2);
+ EXPECT_EQ(frames.at(node_map->GetNode("InvariantAdd2")).size(), 2);
+ EXPECT_EQ(frames.at(node_map->GetNode("InvariantAdd2")).back(), 1);
+ EXPECT_EQ(frames.at(node_map->GetNode("Const2")).size(), 2);
+ EXPECT_EQ(frames.at(node_map->GetNode("Const2")).back(), 1);
+
+ node_map.reset(new NodeMap(&output));
+ EXPECT_TRUE(IdentifyFrames(output, &frames, &num_frames).ok());
+ EXPECT_EQ(num_frames, 2);
+ EXPECT_EQ(frames.at(node_map->GetNode("InvariantAdd2")).size(), 1);
+ EXPECT_EQ(frames.at(node_map->GetNode("InvariantAdd2")).back(), 0);
+ EXPECT_EQ(frames.at(node_map->GetNode("Const2")).size(), 1);
+ EXPECT_EQ(frames.at(node_map->GetNode("Const2")).back(), 0);
+}
+
+TEST_F(LoopOptimizerTest, NestedLoopConst2) {
+ GraphDef graph;
+ *graph.add_node() = CreateNode("0", {});
+ *graph.add_node() = CreateNode(
+ "InvariantEnter", "Enter", "while/while_context", true, 1, {"0"});
+ *graph.add_node() = CreateNode(
+ "InvariantAdd", "Add", {"InvariantEnter", "InvariantEnter"});
+ *graph.add_node() = CreateNode(
+ "VariantAdd", "Add", {"InvariantAdd", "Identity"});
+ *graph.add_node() = CreateNode(
+ "VariantEnter", "Enter", "while/while_context", false, 1, {"0"});
+ *graph.add_node() = CreateNode(
+ "Merge", "Merge", {"VariantEnter", "NextIteration"});
+ *graph.add_node() = CreateNode("Less/y", "Const", {"^Identity"});
+ *graph.add_node() = CreateNode("Less", "Less", {"Exit2", "less/y"});
+ *graph.add_node() = CreateNode("LoopCond", "LoopCond", {"Less"});
+ *graph.add_node() = CreateNode("Switch", "Switch", {"Merge", "LoopCond"});
+ *graph.add_node() = CreateNode("Identity", {"Switch:1"});
+ *graph.add_node() = CreateNode(
+ "NextIteration", "NextIteration", {"Exit2"});
+ *graph.add_node() = CreateNode("Exit", "Exit", {"Switch"});
+ *graph.add_node() = CreateNode("1", {"Exit"});
+
+ *graph.add_node() = CreateNode(
+ "InvariantEnter2", "Enter", "while/while/while_context", true, 1,
+ {"InvariantAdd"});
+ *graph.add_node() = CreateNode("Const2", "Const", {"^Identity2"});
+ *graph.add_node() = CreateNode(
+ "InvariantAdd2", "Add", {"InvariantEnter2", "Const2"});
+ *graph.add_node() = CreateNode(
+ "VariantAdd2", "Add", {"InvariantAdd2", "Identity2"});
+ *graph.add_node() = CreateNode(
+ "VariantEnter2", "Enter", "while/while/while_context", false, 1,
+ {"VariantEnter"});
+ *graph.add_node() = CreateNode(
+ "Merge2", "Merge", {"VariantEnter2", "NextIteration2"});
+ *graph.add_node() = CreateNode("Less2/y", "Const", {"^Identity2"});
+ *graph.add_node() = CreateNode("Less2", "Less", {"VariantAdd2", "less2/y"});
+ *graph.add_node() = CreateNode("LoopCond2", "LoopCond", {"Less2"});
+ *graph.add_node() = CreateNode("Switch2", "Switch", {"Merge2", "LoopCond2"});
+ *graph.add_node() = CreateNode("Identity2", {"Switch2:1"});
+ *graph.add_node() = CreateNode(
+ "NextIteration2", "NextIteration", {"VariantAdd2"});
+ *graph.add_node() = CreateNode("Exit2", "Exit", {"Switch2"});
+
+ GrapplerItem item;
+ item.graph = graph;
+
+ LoopOptimizer optimizer;
+ GraphDef output;
+ Status status = optimizer.Optimize(nullptr, item, &output);
+ TF_EXPECT_OK(status);
+
+ std::unique_ptr<NodeMap> node_map;
+ std::unordered_map<const NodeDef*, std::vector<int>> frames;
+ int num_frames;
+
+ node_map.reset(new NodeMap(&graph));
+ EXPECT_TRUE(IdentifyFrames(graph, &frames, &num_frames).ok());
+ EXPECT_EQ(num_frames, 2);
+ EXPECT_EQ(frames.at(node_map->GetNode("InvariantAdd2")).size(), 2);
+ EXPECT_EQ(frames.at(node_map->GetNode("InvariantAdd2")).back(), 1);
+ EXPECT_EQ(frames.at(node_map->GetNode("Const2")).size(), 2);
+ EXPECT_EQ(frames.at(node_map->GetNode("Const2")).back(), 1);
+
+ node_map.reset(new NodeMap(&output));
+ EXPECT_TRUE(IdentifyFrames(output, &frames, &num_frames).ok());
+ EXPECT_EQ(num_frames, 2);
+ EXPECT_EQ(frames.at(node_map->GetNode("InvariantAdd2")).size(), 0);
+ EXPECT_EQ(frames.at(node_map->GetNode("Const2")).size(), 0);
+}
void VerifyGraphsEqual(const GraphDef& original_graph,
const GraphDef& optimized_graph, const string& func) {
srcs = [
"dequantize_op.cc",
"meta_support.cc",
- "quantization_utils.cc",
"quantize_down_and_shrink_range.cc",
"quantize_op.cc",
"quantized_activation_ops.cc",
":image_resizer_state",
":ops_util",
":pooling_ops",
+ ":quantization_utils",
"//tensorflow/core:array_ops_op_lib",
"//tensorflow/core:core_cpu",
"//tensorflow/core:framework",
)
cc_library(
+ name = "quantization_utils",
+ srcs = ["quantization_utils.cc"],
+ hdrs = ["quantization_utils.h"],
+ deps = [
+ "//tensorflow/core:framework",
+ "@gemmlowp",
+ ],
+)
+
+cc_library(
name = "remote_fused_graph_execute_utils",
srcs = [
"i_remote_fused_graph_ops_definitions.cc",
srcs = [
"cwise_ops_common.cc",
"meta_support.cc",
- "quantization_utils.cc",
],
hdrs = [
"cwise_ops.h",
"cwise_ops_gpu_gradients.cu.h",
"cwise_ops_gradients.h",
"meta_support.h",
- "quantization_utils.h",
],
deps = [
":bounds_check",
+ ":quantization_utils",
"//tensorflow/core:framework",
"//tensorflow/core:lib",
"//third_party/eigen3",
#include "tensorflow/core/kernels/cwise_ops_common.h"
namespace tensorflow {
-REGISTER5(BinaryOp, CPU, "Maximum", functor::maximum, float, Eigen::half,
- double, int32, int64);
+REGISTER6(BinaryOp, CPU, "Maximum", functor::maximum, float, Eigen::half,
+ bfloat16, double, int32, int64);
#if GOOGLE_CUDA
REGISTER4(BinaryOp, GPU, "Maximum", functor::maximum, float, Eigen::half,
double, int64);
return;
}
- if (dnn_shape_src.IsMklTensor())
- depth_ = dnn_shape_src.DimSize(MklDnnDims::Dim_C);
- else
- ExtractParams(context);
-
- memory::format format_m;
if (dnn_shape_src.IsMklTensor()) {
- if (dnn_shape_src.IsTensorInNCHWFormat())
- format_m = memory::format::nchw;
- else
- format_m = memory::format::nhwc;
+ depth_ = dnn_shape_src.DimSize(MklDnnDims::Dim_C);
+ } else if (dnn_shape_diff_dst.IsMklTensor()) {
+ depth_ = dnn_shape_diff_dst.DimSize(MklDnnDims::Dim_C);
} else {
- format_m = TFDataFormatToMklDnnDataFormat(tensor_format_);
+ ExtractParams(context);
}
MklDnnData<T> src(&cpu_engine);
diff_dst_dims =
TFShapeToMklDnnDimsInNCHW(diff_dst_tensor.shape(), tensor_format_);
- // set src and diff_dst primitives
+ // set src and diff_dst primitives according to input layout
memory::desc src_md({}, memory::data_undef, memory::format_undef);
memory::desc diff_dst_md({}, memory::data_undef, memory::format_undef);
- if (dnn_shape_src.IsMklTensor() || dnn_shape_diff_dst.IsMklTensor()) {
- if (dnn_shape_src.IsMklTensor()) {
- src_md = dnn_shape_src.GetMklLayout();
- diff_dst_md = src_md;
- } else {
- diff_dst_md = dnn_shape_diff_dst.GetMklLayout();
- src_md = diff_dst_md;
- }
+ if (dnn_shape_src.IsMklTensor()) {
+ src_md = dnn_shape_src.GetMklLayout();
} else {
- src_md = memory::desc(src_dims, MklDnnType<T>(), format_m);
- diff_dst_md = src_md;
+ src_md = memory::desc(src_dims, MklDnnType<T>(),
+ TFDataFormatToMklDnnDataFormat(tensor_format_));
+ }
+ if (dnn_shape_diff_dst.IsMklTensor()) {
+ diff_dst_md = dnn_shape_diff_dst.GetMklLayout();
+ } else {
+ diff_dst_md = memory::desc(diff_dst_dims, MklDnnType<T>(),
+ TFDataFormatToMklDnnDataFormat(tensor_format_));
}
src.SetUsrMem(src_md, &src_tensor);
diff_dst.SetUsrMem(diff_dst_md, &diff_dst_tensor);
// allocate diff_src tensor
MklDnnShape dnn_shape_diff_src;
TensorShape tf_shape_diff_src;
- if (dnn_shape_src.IsMklTensor()) {
+
+ // MKL-DNN's BN primitive not provide API to fetch internal format
+ // set common_md as OpMem
+ // src and diff_dst will reorder to common_md
+ // diff_src will set as common_md
+ memory::desc common_md({}, memory::data_undef, memory::format_undef);
+ if (dnn_shape_src.IsMklTensor() || dnn_shape_diff_dst.IsMklTensor()) {
+ if (dnn_shape_src.IsMklTensor()) {
+ common_md = dnn_shape_src.GetMklLayout();
+ } else {
+ common_md = dnn_shape_diff_dst.GetMklLayout();
+ }
+ } else {
+ common_md = memory::desc(src_dims, MklDnnType<T>(),
+ TFDataFormatToMklDnnDataFormat(tensor_format_));
+ }
+ // if any of src and diff_dst as mkl layout,
+ // then we set diff_src as mkl layout
+ if (dnn_shape_src.IsMklTensor() ||
+ dnn_shape_diff_dst.IsMklTensor()) {
dnn_shape_diff_src.SetMklTensor(true);
- auto diff_src_pd = bnrm_fwd_pd.dst_primitive_desc();
+ // set diff_src's mkl layout as common_md
+ auto diff_src_pd = memory::primitive_desc(common_md, cpu_engine);
dnn_shape_diff_src.SetMklLayout(&diff_src_pd);
dnn_shape_diff_src.SetElemType(MklDnnType<T>());
- dnn_shape_diff_src.SetTfLayout(dnn_shape_src.GetDimension(), src_dims,
- format_m);
- dnn_shape_diff_src.SetTfDimOrder(dnn_shape_src.GetDimension(),
- tensor_format_);
+ if (dnn_shape_src.IsMklTensor()) {
+ dnn_shape_diff_src.SetTfLayout(
+ dnn_shape_src.GetDimension(),
+ src_dims,
+ dnn_shape_src.GetTfDataFormat());
+ dnn_shape_diff_src.SetTfDimOrder(
+ dnn_shape_src.GetDimension(),
+ tensor_format_);
+ } else {
+ dnn_shape_diff_src.SetTfLayout(
+ dnn_shape_diff_dst.GetDimension(),
+ src_dims,
+ dnn_shape_diff_dst.GetTfDataFormat());
+ dnn_shape_diff_src.SetTfDimOrder(
+ dnn_shape_diff_dst.GetDimension(),
+ tensor_format_);
+ }
tf_shape_diff_src.AddDim(diff_src_pd.get_size() / sizeof(T));
} else {
dnn_shape_diff_src.SetMklTensor(false);
+ // both src and diff_dst are TensorFlow layout,
+ // so it is OK to get TensorFlow shape.
tf_shape_diff_src = src_tensor.shape();
}
AllocateOutputSetMklShape(context, kDiffSrcIndex, &diff_src_tensor,
tf_shape_diff_src, dnn_shape_diff_src);
- diff_src.SetUsrMem(src_md, diff_src_tensor);
+ // set diff_src
+ diff_src.SetUsrMem(common_md, diff_src_tensor);
prop_kind pk = prop_kind::backward;
auto bnrm_bwd_desc = batch_normalization_backward::desc(
- pk, diff_src.GetUsrMemDesc(), src.GetUsrMemDesc(), epsilon_,
+ pk, common_md, common_md, epsilon_,
/* for inference, specify use_global_stats
1. on fwd prop, use mean and variance
provided as inputs
auto bnrm_bwd_pd = batch_normalization_backward::primitive_desc(
bnrm_bwd_desc, cpu_engine, bnrm_fwd_pd);
+ std::vector<primitive> net;
+ src.CheckReorderToOpMem(memory::primitive_desc(common_md,
+ cpu_engine), &net);
+ diff_dst.CheckReorderToOpMem(memory::primitive_desc(common_md,
+ cpu_engine), &net);
+
auto bnrm_bwd_op = batch_normalization_backward(
bnrm_bwd_pd, src.GetOpMem(), mean.GetOpMem(), variance.GetOpMem(),
diff_dst.GetOpMem(), weights_m, diff_src.GetOpMem(), diff_weights_m);
- std::vector<primitive> net;
net.push_back(bnrm_bwd_op);
stream(stream::kind::eager).submit(net).wait();
mkl_context.MklCleanup();
}
+
+
#else // INTEL_MKL_ML
+
template <typename Device, typename T, algorithm alg_kind>
class MklReluOpBase : public OpKernel {
public:
// allocate diff_src tensor
MklDnnShape dnn_shape_diff_src;
TensorShape tf_shape_diff_src;
- if (dnn_shape_src.IsMklTensor()) {
+ if (dnn_shape_src.IsMklTensor() ||
+ dnn_shape_diff_dst.IsMklTensor()) {
dnn_shape_diff_src.SetMklTensor(true);
auto diff_src_pd = relu_bwd_pd.diff_src_primitive_desc();
dnn_shape_diff_src.SetMklLayout(&diff_src_pd);
dnn_shape_diff_src.SetElemType(MklDnnType<T>());
- dnn_shape_diff_src.SetTfLayout(dnn_shape_src.GetDimension(),
- dnn_shape_src.GetSizesAsMklDnnDims(),
- dnn_shape_src.GetTfDataFormat());
+ if (dnn_shape_src.IsMklTensor()) {
+ dnn_shape_diff_src.SetTfLayout(dnn_shape_src.GetDimension(),
+ dnn_shape_src.GetSizesAsMklDnnDims(),
+ dnn_shape_src.GetTfDataFormat());
+ } else {
+ dnn_shape_diff_src.SetTfLayout(dnn_shape_diff_dst.GetDimension(),
+ dnn_shape_diff_dst.GetSizesAsMklDnnDims(),
+ dnn_shape_diff_dst.GetTfDataFormat());
+ }
tf_shape_diff_src.AddDim(diff_src_pd.get_size() / sizeof(T));
} else {
dnn_shape_diff_src.SetMklTensor(false);
+ // both src and diff_dst are TensorFlow layout,
+ // so it is ok to get TensorFlow shape.
tf_shape_diff_src = src_tensor.shape();
}
AllocateOutputSetMklShape(context, diff_src_index, &diff_src_tensor,
.TypeConstraint<int64>("Tshape"), \
ReshapeOp);
TF_CALL_NUMBER_TYPES_NO_INT32(REGISTER_GPU_KERNEL);
-TF_CALL_bfloat16(REGISTER_GPU_KERNEL);
TF_CALL_bool(REGISTER_GPU_KERNEL);
#undef REGISTER_GPU_KERNEL
#define EIGEN_USE_GPU
#endif // GOOGLE_CUDA
-#include "tensorflow/core/kernels/segment_reduction_ops.h"
-#include <vector>
#include "third_party/eigen3/Eigen/Core"
#include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor"
+#include "tensorflow/core/kernels/segment_reduction_ops.h"
+#include <vector>
#include "tensorflow/core/framework/numeric_op.h"
#include "tensorflow/core/framework/op_kernel.h"
#include "tensorflow/core/framework/register_types.h"
#undef REGISTER_GPU_SORTED_KERNELS_ALL
#endif // GOOGLE_CUDA
+// ____________________________________________________________________________
+// Unsorted segment reduction ops.
+
namespace functor {
-// UnsortedSegmentSumFunctor implementation for CPUDevice.
-// todo: Remove duplicate code in UnsortedSegmentSumFunctor and
-// UnsortedSegmentMaxFunctor.
-template <typename T, typename Index>
-struct UnsortedSegmentSumFunctor<CPUDevice, T, Index>
- : UnsortedSegmentBaseFunctor<CPUDevice, T, Index> {
- void operator()(OpKernelContext* ctx, const CPUDevice& d,
- const Index output_rows, const TensorShape& segment_ids_shape,
+// The ReductionFunctor implementation for CPU.
+template <typename T, typename Index, typename InitialValueF,
+ typename ReductionF>
+struct UnsortedSegmentFunctor<CPUDevice, T, Index, InitialValueF, ReductionF> {
+ void operator()(OpKernelContext* ctx, const Index num_segments,
+ const TensorShape& segment_ids_shape,
typename TTypes<Index>::ConstFlat segment_ids,
const Index data_size, const T* data,
- typename TTypes<T, 2>::Tensor output) override {
- output.setZero();
+ typename TTypes<T, 2>::Tensor output) {
+ output.setConstant(InitialValueF()());
if (data_size == 0) {
return;
}
const int64 N = segment_ids.dimension(0);
+ ReductionF reduction;
auto data_flat = typename TTypes<T, 2>::ConstTensor(data, N, data_size / N);
for (int64 i = 0; i < N; ++i) {
Index j = internal::SubtleMustCopy(segment_ids(i));
if (j < 0) {
continue;
}
- OP_REQUIRES(ctx, FastBoundsCheck(j, output_rows),
+ OP_REQUIRES(ctx, FastBoundsCheck(j, num_segments),
errors::InvalidArgument(
"segment_ids", SliceDebugString(segment_ids_shape, i),
- " = ", j, " is out of range [0, ", output_rows, ")"));
- output.template chip<0>(j) += data_flat.template chip<0>(i);
+ " = ", j, " is out of range [0, ", num_segments, ")"));
+ reduction(data_flat.template chip<0>(i), output.template chip<0>(j));
}
}
};
-// UnsortedSegmentMaxFunctor implementation for CPUDevice.
-template <typename T, typename Index>
-struct UnsortedSegmentMaxFunctor<CPUDevice, T, Index>
- : UnsortedSegmentBaseFunctor<CPUDevice, T, Index> {
- void operator()(OpKernelContext* ctx, const CPUDevice& d,
- const Index output_rows, const TensorShape& segment_ids_shape,
- typename TTypes<Index>::ConstFlat segment_ids,
- const Index data_size, const T* data,
- typename TTypes<T, 2>::Tensor output) override {
- output.setConstant(std::numeric_limits<T>::lowest());
- if (data_size == 0) {
- return;
- }
- const int64 N = segment_ids.dimension(0);
- auto data_flat = typename TTypes<T, 2>::ConstTensor(data, N, data_size / N);
- for (int64 i = 0; i < N; ++i) {
- Index j = internal::SubtleMustCopy(segment_ids(i));
- OP_REQUIRES(ctx, FastBoundsCheck(j, output_rows),
- errors::InvalidArgument(
- "segment_ids", SliceDebugString(segment_ids_shape, i),
- " = ", j, " is out of range [0, ", output_rows, ")"));
- output.template chip<0>(j) =
- data_flat.template chip<0>(i).cwiseMax(output.template chip<0>(j));
- }
+
+template <typename T>
+using MatrixChip = Eigen::TensorChippingOp<0l, typename TTypes<T, 2>::Matrix>;
+
+template <typename T>
+using constMatrixChip =
+ Eigen::TensorChippingOp<0l, const typename TTypes<T, 2>::ConstMatrix>;
+
+// reduction functors
+template <typename T>
+struct SumOp {
+ void operator()(const constMatrixChip<T> data, MatrixChip<T> output) {
+ output += data;
+ }
+};
+
+template <typename T>
+struct MaxOp {
+ void operator()(const constMatrixChip<T> data, MatrixChip<T> output) {
+ output = data.cwiseMax(output);
+ }
+};
+
+template <typename T>
+struct MinOp {
+ void operator()(const constMatrixChip<T> data, MatrixChip<T> output) {
+ output = data.cwiseMin(output);
+ }
+};
+
+template <typename T>
+struct ProdOp {
+ void operator()(const constMatrixChip<T> data, MatrixChip<T> output) {
+ output *= data;
}
};
} // namespace functor
-// Base class for SegmentReductionOps that can handle unsorted segment
-// definitions
-// and specifying the size of the output in addition to a reduction function
-template <typename Device, class T, class Index>
-class UnsortedSegmentBaseOp : public OpKernel {
+// Static check routines not in the templated class to reduce code size
+static void UnsortedSegmentReductionValidation(OpKernel* op_kernel,
+ OpKernelContext* context,
+ const Tensor& data,
+ const Tensor& segment_ids,
+ const Tensor& num_segments) {
+ OP_REQUIRES(
+ context, op_kernel->IsLegacyScalar(num_segments.shape()),
+ errors::InvalidArgument("num_segments should be a scalar, not shape ",
+ num_segments.shape().DebugString()));
+ OP_REQUIRES(
+ context, TensorShapeUtils::StartsWith(data.shape(), segment_ids.shape()),
+ errors::InvalidArgument("data.shape = ", data.shape().DebugString(),
+ " does not start with segment_ids.shape = ",
+ segment_ids.shape().DebugString()));
+}
+
+static bool UnsortedSegmentReductionDoValidation(OpKernel* op_kernel,
+ OpKernelContext* context,
+ const Tensor& data,
+ const Tensor& segment_ids,
+ const Tensor& num_segments) {
+ UnsortedSegmentReductionValidation(op_kernel, context, data, segment_ids,
+ num_segments);
+ return context->status().ok();
+}
+
+// The UnsortedSegmentReduction OpKernel. The DeviceReductionFunctor
+// is the device specific implementation of the reduction. These device
+// specific implementations are templated themselves with the corresponding
+// initial value functors and reduction functors.
+template <typename T, typename Index, typename DeviceReductionFunctor>
+class UnsortedSegmentReductionOp : public OpKernel {
public:
- explicit UnsortedSegmentBaseOp(
- OpKernelConstruction* context,
- functor::UnsortedSegmentBaseFunctor<Device, T, Index>& functor)
- : OpKernel(context), reduction_functor_(functor) {}
+ explicit UnsortedSegmentReductionOp(OpKernelConstruction* context)
+ : OpKernel(context), reduction_functor_(DeviceReductionFunctor()) {}
void Compute(OpKernelContext* context) override {
const Tensor& data = context->input(0);
const Tensor& segment_ids = context->input(1);
const Tensor& num_segments = context->input(2);
-
- OP_REQUIRES(
- context, IsLegacyScalar(num_segments.shape()),
- errors::InvalidArgument("num_segments should be a scalar, not shape ",
- num_segments.shape().DebugString()));
- OP_REQUIRES(
- context,
- TensorShapeUtils::StartsWith(data.shape(), segment_ids.shape()),
- errors::InvalidArgument("data.shape = ", data.shape().DebugString(),
- " does not start with segment_ids.shape = ",
- segment_ids.shape().DebugString()));
-
+ if (!UnsortedSegmentReductionDoValidation(this, context, data, segment_ids,
+ num_segments)) {
+ return;
+ }
const auto segment_flat = segment_ids.flat<Index>();
const Index output_rows =
internal::SubtleMustCopy(num_segments.scalar<int32>()());
OP_REQUIRES(context, output_rows >= 0,
errors::InvalidArgument("Input num_segments == ", output_rows,
" must not be negative."));
-
TensorShape output_shape;
output_shape.AddDim(output_rows);
for (int i = segment_ids.dims(); i < data.dims(); i++) {
output_shape.AddDim(data.dim_size(i));
}
-
Tensor* output = nullptr;
OP_REQUIRES_OK(context, context->allocate_output(0, output_shape, &output));
auto output_flat = output->flat_outer_dims<T>();
-
auto data_ptr = data.template flat<T>().data();
- reduction_functor_(context, context->template eigen_device<Device>(),
- output_rows, segment_ids.shape(), segment_flat,
+ reduction_functor_(context, output_rows, segment_ids.shape(), segment_flat,
data.NumElements(), data_ptr, output_flat);
}
- private:
- functor::UnsortedSegmentBaseFunctor<Device, T, Index>& reduction_functor_;
-};
-
-template <typename Device, class T, class Index>
-class UnsortedSegmentSumOp : public UnsortedSegmentBaseOp<Device, T, Index> {
- public:
- explicit UnsortedSegmentSumOp(OpKernelConstruction* context)
- : UnsortedSegmentBaseOp<Device, T, Index>(context, sum_functor_) {}
-
- private:
- functor::UnsortedSegmentSumFunctor<Device, T, Index> sum_functor_;
+ protected:
+ DeviceReductionFunctor reduction_functor_;
};
-template <typename Device, class T, class Index>
-class UnsortedSegmentMaxOp : public UnsortedSegmentBaseOp<Device, T, Index> {
- public:
- explicit UnsortedSegmentMaxOp(OpKernelConstruction* context)
- : UnsortedSegmentBaseOp<Device, T, Index>(context, max_functor_) {}
-
- private:
- functor::UnsortedSegmentMaxFunctor<Device, T, Index> max_functor_;
-};
-
-#define REGISTER_REAL_CPU_UNSORTED_KERNELS(type, index_type) \
- REGISTER_KERNEL_BUILDER(Name("UnsortedSegmentSum") \
- .Device(DEVICE_CPU) \
- .TypeConstraint<type>("T") \
- .TypeConstraint<index_type>("Tindices"), \
- UnsortedSegmentSumOp<CPUDevice, type, index_type>); \
- REGISTER_KERNEL_BUILDER(Name("UnsortedSegmentMax") \
- .Device(DEVICE_CPU) \
- .TypeConstraint<type>("T") \
- .TypeConstraint<index_type>("Tindices"), \
- UnsortedSegmentMaxOp<CPUDevice, type, index_type>);
-
-#define REGISTER_COMPLEX_CPU_UNSORTED_KERNELS(type, index_type) \
- REGISTER_KERNEL_BUILDER(Name("UnsortedSegmentSum") \
- .Device(DEVICE_CPU) \
- .TypeConstraint<type>("T") \
- .TypeConstraint<index_type>("Tindices"), \
- UnsortedSegmentSumOp<CPUDevice, type, index_type>);
+#define REGISTER_CPU_KERNEL_UNSORTEDSEGMENT( \
+ name, type, index_type, initial_value_functor, reduction_functor) \
+ REGISTER_KERNEL_BUILDER( \
+ Name(name) \
+ .Device(DEVICE_CPU) \
+ .TypeConstraint<type>("T") \
+ .TypeConstraint<index_type>("Tindices"), \
+ UnsortedSegmentReductionOp< \
+ type, index_type, \
+ functor::UnsortedSegmentFunctor<CPUDevice, type, index_type, \
+ initial_value_functor, \
+ reduction_functor> >)
+
+#define REGISTER_REAL_CPU_UNSORTED_KERNELS(type, index_type) \
+ REGISTER_CPU_KERNEL_UNSORTEDSEGMENT("UnsortedSegmentSum", type, index_type, \
+ functor::Zero<type>, \
+ functor::SumOp<type>); \
+ REGISTER_CPU_KERNEL_UNSORTEDSEGMENT("UnsortedSegmentMax", type, index_type, \
+ functor::Lowest<type>, \
+ functor::MaxOp<type>); \
+ REGISTER_CPU_KERNEL_UNSORTEDSEGMENT("UnsortedSegmentMin", type, index_type, \
+ functor::Highest<type>, \
+ functor::MinOp<type>); \
+ REGISTER_CPU_KERNEL_UNSORTEDSEGMENT("UnsortedSegmentProd", type, index_type, \
+ functor::One<type>, \
+ functor::ProdOp<type>);
+
+#define REGISTER_COMPLEX_CPU_UNSORTED_KERNELS(type, index_type) \
+ REGISTER_CPU_KERNEL_UNSORTEDSEGMENT("UnsortedSegmentSum", type, index_type, \
+ functor::Zero<type>, \
+ functor::SumOp<type>); \
+ REGISTER_CPU_KERNEL_UNSORTEDSEGMENT("UnsortedSegmentProd", type, index_type, \
+ functor::One<type>, \
+ functor::ProdOp<type>)
#define REGISTER_REAL_CPU_UNSORTED_KERNELS_ALL(type) \
REGISTER_REAL_CPU_UNSORTED_KERNELS(type, int32); \
TF_CALL_REAL_NUMBER_TYPES(REGISTER_REAL_CPU_UNSORTED_KERNELS_ALL);
REGISTER_COMPLEX_CPU_UNSORTED_KERNELS_ALL(complex64);
REGISTER_COMPLEX_CPU_UNSORTED_KERNELS_ALL(complex128);
+
#undef REGISTER_REAL_CPU_UNSORTED_KERNELS
+#undef REGISTER_CPU_KERNEL_UNSORTEDSEGMENT
#undef REGISTER_COMPLEX_CPU_UNSORTED_KERNELS
#undef REGISTER_COMPLEX_CPU_UNSORTED_KERNELS_ALL
#undef REGISTER_REAL_CPU_UNSORTED_KERNELS_ALL
#if GOOGLE_CUDA
-#define REGISTER_GPU_UNSORTED_KERNELS(type, index_type) \
- REGISTER_KERNEL_BUILDER(Name("UnsortedSegmentSum") \
- .Device(DEVICE_GPU) \
- .HostMemory("num_segments") \
- .TypeConstraint<type>("T") \
- .TypeConstraint<index_type>("Tindices"), \
- UnsortedSegmentSumOp<GPUDevice, type, index_type>);
-
-#define REGISTER_GPU_UNSORTED_KERNELS_ALL(type) \
- REGISTER_GPU_UNSORTED_KERNELS(type, int32); \
- REGISTER_GPU_UNSORTED_KERNELS(type, int64);
+#define REGISTER_GPU_KERNEL_UNSORTEDSEGMENT( \
+ name, type, index_type, initial_value_functor, reduction_kernel_functor) \
+ REGISTER_KERNEL_BUILDER( \
+ Name(name) \
+ .Device(DEVICE_GPU) \
+ .HostMemory("num_segments") \
+ .TypeConstraint<type>("T") \
+ .TypeConstraint<index_type>("Tindices"), \
+ UnsortedSegmentReductionOp< \
+ type, index_type, \
+ functor::UnsortedSegmentFunctor<GPUDevice, type, index_type, \
+ initial_value_functor, \
+ reduction_kernel_functor> >)
+
+// sum is the only op that supports all input types currently
+#define REGISTER_REAL_GPU_UNSORTED_KERNELS(type, index_type) \
+ REGISTER_GPU_KERNEL_UNSORTEDSEGMENT("UnsortedSegmentMax", type, index_type, \
+ functor::Lowest<type>, \
+ functor::MaxOpGpu<type>); \
+ REGISTER_GPU_KERNEL_UNSORTEDSEGMENT("UnsortedSegmentMin", type, index_type, \
+ functor::Highest<type>, \
+ functor::MinOpGpu<type>); \
+ REGISTER_GPU_KERNEL_UNSORTEDSEGMENT("UnsortedSegmentProd", type, index_type, \
+ functor::One<type>, \
+ functor::ProdOpGpu<type>);
+
+#define REGISTER_SUM_GPU_UNSORTED_KERNELS(type, index_type) \
+ REGISTER_GPU_KERNEL_UNSORTEDSEGMENT("UnsortedSegmentSum", type, index_type, \
+ functor::Zero<type>, \
+ functor::SumOpGpu<type>);
+
+#define REGISTER_REAL_GPU_UNSORTED_KERNELS_ALL(type) \
+ REGISTER_REAL_GPU_UNSORTED_KERNELS(type, int32); \
+ REGISTER_REAL_GPU_UNSORTED_KERNELS(type, int64);
+
+#define REGISTER_SUM_GPU_UNSORTED_KERNELS_ALL(type) \
+ REGISTER_SUM_GPU_UNSORTED_KERNELS(type, int32); \
+ REGISTER_SUM_GPU_UNSORTED_KERNELS(type, int64);
+
+
+TF_CALL_GPU_NUMBER_TYPES(REGISTER_REAL_GPU_UNSORTED_KERNELS_ALL);
+TF_CALL_int32(REGISTER_REAL_GPU_UNSORTED_KERNELS_ALL);
+TF_CALL_GPU_NUMBER_TYPES(REGISTER_SUM_GPU_UNSORTED_KERNELS_ALL);
+TF_CALL_int32(REGISTER_SUM_GPU_UNSORTED_KERNELS_ALL);
+TF_CALL_complex64(REGISTER_SUM_GPU_UNSORTED_KERNELS_ALL);
+TF_CALL_complex128(REGISTER_SUM_GPU_UNSORTED_KERNELS_ALL);
+
+#undef REGISTER_GPU_KERNEL_UNSORTEDSEGMENT
+#undef REGISTER_REAL_GPU_UNSORTED_KERNELS
+#undef REGISTER_SUM_GPU_UNSORTED_KERNELS
+#undef REGISTER_REAL_GPU_UNSORTED_KERNELS_ALL
+#undef REGISTER_SUM_GPU_UNSORTED_KERNELS_ALL
-TF_CALL_GPU_NUMBER_TYPES(REGISTER_GPU_UNSORTED_KERNELS_ALL);
-TF_CALL_complex64(REGISTER_GPU_UNSORTED_KERNELS_ALL);
-TF_CALL_complex128(REGISTER_GPU_UNSORTED_KERNELS_ALL);
-#undef REGISTER_GPU_UNSORTED_KERNELS
-#undef REGISTER_GPU_UNSORTED_KERNELS_ALL
#endif // GOOGLE_CUDA
+// ____________________________________________________________________________
+// Sparse segment reduction ops.
+
// Same as SegmentReductionOp but takes as input a "sparse" tensor, represented
// by two dense tensors, one containing the data, and the other containing
// indices into the data.
#ifndef TENSORFLOW_CORE_KERNELS_SEGMENT_REDUCTION_OPS_H_
#define TENSORFLOW_CORE_KERNELS_SEGMENT_REDUCTION_OPS_H_
+// This file requires the following include because it uses CudaAtomicMax:
+// #include "tensorflow/core/util/cuda_kernel_helper.h"
+
+// Unfortunately we can't add the #include, since it breaks compilation for
+// non-GPU targets. This only breaks in clang, because it's more strict for
+// template code and CudaAtomicMax is used in template context.
+
#include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor"
#include "tensorflow/core/framework/tensor.h"
#include "tensorflow/core/framework/tensor_shape.h"
const Index data_size, const T* data,
typename TTypes<T, 2>::Tensor output);
};
-#endif
-// BaseFunctor for definition of UnsorteSegmentReductionOp
-// for usage without templates.
-template <typename Device, typename T, typename Index>
-struct UnsortedSegmentBaseFunctor {
- virtual ~UnsortedSegmentBaseFunctor() {}
- virtual void operator()(OpKernelContext* ctx, const Device& d,
- const Index output_rows,
- const TensorShape& segment_ids_shape,
- typename TTypes<Index>::ConstFlat segment_ids,
- const Index data_size, const T* data,
- typename TTypes<T, 2>::Tensor output){};
-};
+#endif
-// Functor for UnsortedSegmentSumOp.
-// output_rows: the number of output segments (unique segment ids in
-// 'segment_ids').
-// segment_ids_shape: shape of 'segment_ids' tensor.
-// segment_ids: unsorted map from input to output segment ids at which to
-// perform segment sum operation.
-// data_size: size of input data tensor.
-// data: input data tensor.
-// output: output reshaped to {output_rows, output.size/output_rows}
-template <typename Device, typename T, typename Index>
-struct UnsortedSegmentSumFunctor
- : public UnsortedSegmentBaseFunctor<Device, T, Index> {
- void operator()(OpKernelContext* ctx, const Device& d,
- const Index output_rows, const TensorShape& segment_ids_shape,
+template <typename Device, typename T, typename Index, typename InitialValueF,
+ typename ReductionF>
+struct UnsortedSegmentFunctor {
+ void operator()(OpKernelContext* ctx, const Index num_segments,
+ const TensorShape& segment_ids_shape,
typename TTypes<Index>::ConstFlat segment_ids,
const Index data_size, const T* data,
typename TTypes<T, 2>::Tensor output);
};
-// Functor for UnsortedSegmentMaxOp.
-// output_rows: the number of output segments (unique segment ids in
-// 'segment_ids').
-// segment_ids_shape: shape of 'segment_ids' tensor.
-// segment_ids: unsorted map from input to output segment ids at which to
-// perform segment sum operation.
-// data_size: size of input data tensor.
-// data: input data tensor.
-// output: output reshaped to {output_rows, output.size/output_rows}
-template <typename Device, typename T, typename Index>
-struct UnsortedSegmentMaxFunctor
- : public UnsortedSegmentBaseFunctor<Device, T, Index> {
- void operator()(OpKernelContext* ctx, const Device& d,
- const Index output_rows, const TensorShape& segment_ids_shape,
- typename TTypes<Index>::ConstFlat segment_ids,
- const Index data_size, const T* data,
- typename TTypes<T, 2>::Tensor output);
+#ifdef GOOGLE_CUDA
+// reduction functors for the gpu
+template <typename T>
+struct SumOpGpu {
+ EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void operator()(T* dest,
+ const T& value) {
+ CudaAtomicAdd(dest, value);
+ }
+};
+
+template <typename T>
+struct ProdOpGpu {
+ EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void operator()(T* dest,
+ const T& value) {
+ CudaAtomicMul(dest, value);
+ }
+};
+
+template <typename T>
+struct MaxOpGpu {
+ EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void operator()(T* dest,
+ const T& value) {
+ CudaAtomicMax(dest, value);
+ }
+};
+
+template <typename T>
+struct MinOpGpu {
+ EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE void operator()(T* dest,
+ const T& value) {
+ CudaAtomicMin(dest, value);
+ }
};
+
+#endif // GOOGLE_CUDA
+
+// initial value functors
+template <typename T>
+struct Zero {
+ EIGEN_STRONG_INLINE T operator()() const { return T(0); }
+};
+
+template <typename T>
+struct One {
+ EIGEN_STRONG_INLINE T operator()() const { return T(1); }
+};
+
+template <typename T>
+struct Lowest {
+ EIGEN_STRONG_INLINE T operator()() const {
+ return Eigen::NumTraits<T>::lowest();
+ }
+};
+
+template <typename T>
+struct Highest {
+ EIGEN_STRONG_INLINE T operator()() const {
+ return Eigen::NumTraits<T>::highest();
+ }
+};
+
} // namespace functor
} // namespace tensorflow
#define EIGEN_USE_GPU
-#include "tensorflow/core/kernels/segment_reduction_ops.h"
+// We need to include cuda_kernel_helper.h before segment_reduction_ops.h
+// See comment in segment_reduction_ops.h for more details.
+#include "tensorflow/core/util/cuda_kernel_helper.h"
+#include "tensorflow/core/kernels/segment_reduction_ops.h"
#include "tensorflow/core/framework/register_types.h"
-#include "tensorflow/core/util/cuda_kernel_helper.h"
+#include "tensorflow/core/util/cuda_device_functions.h"
+
namespace tensorflow {
using GPUDevice = Eigen::GpuDevice;
-// Helper for UnusortedSegmentSumCustomKernel that adds value into dest
-// atomically.
-template <typename T>
-static __device__ __forceinline__ void AccumulateInto(T* dest, const T& value) {
- CudaAtomicAdd(dest, value);
-}
-
-// Specializations of AccumulateInto for complex types, which CudaAtomicAdd does
-// not support. We treat a std::complex<T>* as a T* (the C++ standard section
-// 26.4.4 allows this explicitly) and atomic add the real and imaginary
-// components individually. The operation as a whole is not atomic, but we can
-// safely treat the components independently for the purpose of accumulating.
-template <>
-__device__ __forceinline__ void AccumulateInto(
- std::complex<float>* dest, const std::complex<float>& value) {
- auto dest_scalar = reinterpret_cast<float*>(dest);
- CudaAtomicAdd(dest_scalar, value.real());
- CudaAtomicAdd(dest_scalar + 1, value.imag());
-}
-
-template <>
-__device__ __forceinline__ void AccumulateInto(
- std::complex<double>* dest, const std::complex<double>& value) {
- auto dest_scalar = reinterpret_cast<double*>(dest);
- CudaAtomicAdd(dest_scalar, value.real());
- CudaAtomicAdd(dest_scalar + 1, value.imag());
-}
-
// SortedSegmentSumFunctor kernel reduces input data just as
// UnsortedSegmentSumCustomKernel does except that input data
// is partitioned along the outer reduction dimension. This is
const Index* segment_ids,
const T* input, T* output,
const Index total_stripe_count) {
- CUDA_1D_KERNEL_LOOP(stripe_index, total_stripe_count) {
+ for (int stripe_index : CudaGridRangeX(total_stripe_count)) {
const Index segment_offset = stripe_index % inner_dim_size;
const Index input_outer_dim_index_base =
stripe_index / inner_dim_size * Index(OuterDimTileSize);
// decide whether to write result to global memory using atomic
// operations
if (last_output_segment_id == first_segment_id) {
- AccumulateInto<T>(output + output_index, sum);
+ CudaAtomicAdd(output + output_index, sum);
} else {
*(output + output_index) = sum;
}
// the following strip.
const Index output_index =
last_output_segment_id * inner_dim_size + segment_offset;
- AccumulateInto<T>(output + output_index, sum);
+ CudaAtomicAdd(output + output_index, sum);
}
}
-// UnsortedSegmentSumFunctor kernel processes 'input_total_size' elements.
+// UnsortedSegmentSumKernel processes 'input_total_size' elements.
// Each element is mapped from input to output by a combination of its
// 'segment_ids' mapping and 'inner_dim_size'.
-template <typename T, typename Index>
-__global__ void UnsortedSegmentSumCustomKernel(
- const Index input_outer_dim_size, const Index inner_dim_size,
- const Index output_outer_dim_size, const Index* segment_ids, const T* input,
- T* output) {
+template <typename T, typename Index, typename KernelReductionFunctor>
+__global__ void UnsortedSegmentCustomKernel(const Index input_outer_dim_size,
+ const Index inner_dim_size,
+ const Index output_outer_dim_size,
+ const Index* segment_ids,
+ const T* input, T* output) {
const Index input_total_size = input_outer_dim_size * inner_dim_size;
const Index output_total_size = output_outer_dim_size * inner_dim_size;
- CUDA_1D_KERNEL_LOOP(input_index, input_total_size) {
+ for (int input_index : CudaGridRangeX(input_total_size)) {
const Index input_segment_index = input_index / inner_dim_size;
const Index segment_offset = input_index % inner_dim_size;
const Index output_segment_index = segment_ids[input_segment_index];
-
if (output_segment_index < 0 || output_segment_index >= output_total_size) {
continue;
}
const Index output_index =
output_segment_index * inner_dim_size + segment_offset;
- AccumulateInto<T>(output + output_index, ldg(input + input_index));
+ KernelReductionFunctor()(output + output_index, ldg(input + input_index));
}
}
<<<config.block_count, config.thread_per_block, 0, d.stream()>>>(
input_outer_dim_size, input_inner_dim_size, output_rows,
segment_ids.data(), data, output.data(), total_stripe_count);
-};
+}
-// UnsortedSegmentSumFunctor implementation for GPUDevice.
-template <typename T, typename Index>
-struct UnsortedSegmentSumFunctor<GPUDevice, T, Index>
- : UnsortedSegmentBaseFunctor<GPUDevice, T, Index> {
- void operator()(OpKernelContext* ctx, const GPUDevice& d,
- const Index output_rows, const TensorShape& segment_ids_shape,
+template <typename T, typename Index, typename InitialValueF,
+ typename ReductionF>
+struct UnsortedSegmentFunctor<GPUDevice, T, Index, InitialValueF, ReductionF> {
+ void operator()(OpKernelContext* ctx, const Index num_segments,
+ const TensorShape& segment_ids_shape,
typename TTypes<Index>::ConstFlat segment_ids,
const Index data_size, const T* data,
- typename TTypes<T, 2>::Tensor output) override {
+ typename TTypes<T, 2>::Tensor output) {
if (output.size() == 0) {
return;
}
- // Set 'output' to zeros.
+ // Set 'output' to initial value.
+ GPUDevice d = ctx->template eigen_device<GPUDevice>();
CudaLaunchConfig config = GetCudaLaunchConfig(output.size(), d);
- SetZero<<<config.block_count, config.thread_per_block, 0, d.stream()>>>(
- output.size(), output.data());
+ SetToValue<<<config.block_count, config.thread_per_block, 0, d.stream()>>>(
+ output.size(), output.data(), InitialValueF()());
if (data_size == 0 || segment_ids_shape.num_elements() == 0) {
return;
}
-
- // Launch kernel to compute unsorted segment sum.
+ // Launch kernel to compute unsorted segment reduction.
// Notes:
- // *) 'input_total_size' is the total number of elements to process.
+ // *) 'data_size' is the total number of elements to process.
// *) 'segment_ids.shape' is a prefix of data's shape.
// *) 'input_outer_dim_size' is the total number of segments to process.
- const Index input_total_size = data_size;
const Index input_outer_dim_size = segment_ids.dimension(0);
- const Index input_inner_dim_size = input_total_size / input_outer_dim_size;
+ const Index input_inner_dim_size = data_size / input_outer_dim_size;
+ config = GetCudaLaunchConfig(data_size, d);
- config = GetCudaLaunchConfig(input_total_size, d);
- UnsortedSegmentSumCustomKernel<T, Index>
+ UnsortedSegmentCustomKernel<T, Index, ReductionF>
<<<config.block_count, config.thread_per_block, 0, d.stream()>>>(
- input_outer_dim_size, input_inner_dim_size, output_rows,
+ input_outer_dim_size, input_inner_dim_size, num_segments,
segment_ids.data(), data, output.data());
}
};
TF_CALL_GPU_NUMBER_TYPES(DEFINE_SORTED_GPU_SPECS);
-#define DEFINE_GPU_SPECS_INDEX(T, Index) \
- template struct UnsortedSegmentSumFunctor<GPUDevice, T, Index>
-
-#define DEFINE_GPU_SPECS(T) \
- DEFINE_GPU_SPECS_INDEX(T, int32); \
- DEFINE_GPU_SPECS_INDEX(T, int64);
-
-TF_CALL_GPU_NUMBER_TYPES(DEFINE_GPU_SPECS);
-TF_CALL_complex64(DEFINE_GPU_SPECS);
-TF_CALL_complex128(DEFINE_GPU_SPECS);
-
-#undef DEFINE_GPU_SPECS
-#undef DEFINE_GPU_SPECS_INDEX
+#define DEFINE_REAL_UNSORTED_GPU_SPECS_INDEX(T, Index) \
+ template struct UnsortedSegmentFunctor< \
+ GPUDevice, T, Index, functor::Lowest<T>, functor::MaxOpGpu<T>>; \
+ template struct UnsortedSegmentFunctor< \
+ GPUDevice, T, Index, functor::Highest<T>, functor::MinOpGpu<T>>; \
+ template struct UnsortedSegmentFunctor<GPUDevice, T, Index, functor::One<T>, \
+ functor::ProdOpGpu<T>>;
+
+// sum is the only op that supports all input types currently
+#define DEFINE_SUM_UNSORTED_GPU_SPECS_INDEX(T, Index) \
+ template struct UnsortedSegmentFunctor< \
+ GPUDevice, T, Index, functor::Zero<T>, functor::SumOpGpu<T>>;
+
+#define DEFINE_REAL_GPU_SPECS(T) \
+ DEFINE_REAL_UNSORTED_GPU_SPECS_INDEX(T, int32); \
+ DEFINE_REAL_UNSORTED_GPU_SPECS_INDEX(T, int64);
+
+#define DEFINE_SUM_GPU_SPECS(T) \
+ DEFINE_SUM_UNSORTED_GPU_SPECS_INDEX(T, int32); \
+ DEFINE_SUM_UNSORTED_GPU_SPECS_INDEX(T, int64);
+
+TF_CALL_GPU_NUMBER_TYPES(DEFINE_REAL_GPU_SPECS);
+TF_CALL_int32(DEFINE_REAL_GPU_SPECS);
+TF_CALL_GPU_NUMBER_TYPES(DEFINE_SUM_GPU_SPECS);
+TF_CALL_int32(DEFINE_SUM_GPU_SPECS);
+TF_CALL_complex64(DEFINE_SUM_GPU_SPECS);
+TF_CALL_complex128(DEFINE_SUM_GPU_SPECS);
+
+#undef DEFINE_SORTED_GPU_SPECS_INDEX
+#undef DEFINE_SORTED_GPU_SPECS
+#undef DEFINE_REAL_UNSORTED_GPU_SPECS_INDEX
+#undef DEFINE_SUM_UNSORTED_GPU_SPECS_INDEX
+#undef DEFINE_REAL_GPU_SPECS
+#undef DEFINE_SUM_GPU_SPECS
} // namespace functor
} // namespace tensorflow
.Device(DEVICE_CPU) \
.TypeConstraint<type>("T") \
.TypeConstraint<int64>("out_idx"), \
+ UniqueOp<type, int64>); \
+ REGISTER_KERNEL_BUILDER(Name("UniqueWithCountsV2") \
+ .Device(DEVICE_CPU) \
+ .TypeConstraint<type>("T") \
+ .TypeConstraint<int32>("out_idx"), \
+ UniqueOp<type, int32>) \
+ REGISTER_KERNEL_BUILDER(Name("UniqueWithCountsV2") \
+ .Device(DEVICE_CPU) \
+ .TypeConstraint<type>("T") \
+ .TypeConstraint<int64>("out_idx"), \
UniqueOp<type, int64>)
TF_CALL_REAL_NUMBER_TYPES(REGISTER_UNIQUE);
REGISTER_UNIQUE(string)
#define EIGEN_USE_THREADS
-#include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor"
#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 "third_party/eigen3/unsupported/Eigen/CXX11/Tensor"
namespace tensorflow {
return Status::OK();
});
+REGISTER_OP("UniqueWithCountsV2")
+ .Input("x: T")
+ .Input("axis: Taxis")
+ .Output("y: T")
+ .Output("idx: out_idx")
+ .Output("count: out_idx")
+ .Attr("T: type")
+ .Attr("Taxis: {int32,int64} = DT_INT64")
+ .Attr("out_idx: {int32, int64} = DT_INT32")
+ .SetShapeFn([](InferenceContext* c) {
+ auto uniq = c->Vector(InferenceContext::kUnknownDim);
+ c->set_output(0, uniq);
+ c->set_output(1, c->input(0));
+ c->set_output(2, uniq);
+ return Status::OK();
+ });
+
namespace {
Status ShapeShapeFn(InferenceContext* c) {
.Attr("Tnumsegments: {int32,int64} = DT_INT32")
.SetShapeFn(UnsortedSegmentReductionShapeFn);
+REGISTER_OP("UnsortedSegmentMin")
+ .Input("data: T")
+ .Input("segment_ids: Tindices")
+ .Input("num_segments: Tnumsegments")
+ .Output("output: T")
+ .Attr("T: realnumbertype")
+ .Attr("Tindices: {int32,int64}")
+ .Attr("Tnumsegments: {int32,int64} = DT_INT32")
+ .SetShapeFn(UnsortedSegmentReductionShapeFn);
+
+REGISTER_OP("UnsortedSegmentProd")
+ .Input("data: T")
+ .Input("segment_ids: Tindices")
+ .Input("num_segments: Tnumsegments")
+ .Output("output: T")
+ .Attr("T: realnumbertype")
+ .Attr("Tindices: {int32,int64}")
+ .Attr("Tnumsegments: {int32,int64} = DT_INT32")
+ .SetShapeFn(UnsortedSegmentReductionShapeFn);
+
REGISTER_OP("SparseSegmentSum")
.Input("data: T")
.Input("indices: Tidx")
#include <aws/core/utils/StringUtils.h>
#include <aws/core/utils/logging/AWSLogging.h>
#include <aws/core/utils/logging/LogSystemInterface.h>
+#include <aws/core/utils/StringUtils.h>
#include <aws/s3/S3Client.h>
#include <aws/s3/S3Errors.h>
#include <aws/s3/model/CopyObjectRequest.h>
#endif
#include <Windows.h>
+#include <shlwapi.h>
#include "tensorflow/core/platform/cpu_info.h"
#include "tensorflow/core/platform/demangle.h"
string Demangle(const char* mangled) { return mangled; }
double NominalCPUFrequency() {
-#ifdef TENSORFLOW_USE_ABSL
- return absl::base_internal::NominalCPUFrequency();
-#else
+ DWORD data;
+ DWORD data_size = sizeof(data);
+ #pragma comment(lib, "shlwapi.lib") // For SHGetValue().
+ if (SUCCEEDED(
+ SHGetValueA(HKEY_LOCAL_MACHINE,
+ "HARDWARE\\DESCRIPTION\\System\\CentralProcessor\\0",
+ "~MHz", nullptr, &data, &data_size))) {
+ return data * 1e6; // Value is MHz.
+ }
return 1.0;
-#endif
}
int64 AvailableRam() {
#include <algorithm>
#include <complex>
+#include "third_party/eigen3/unsupported/Eigen/CXX11/Tensor"
#include "cuda/include/cuda.h"
-#include "cuda/include/device_functions.h"
#include "tensorflow/core/platform/types.h"
-#if CUDA_VERSION >= 7050
-#include "cuda/include/cuda_fp16.h"
-#endif // CUDA_VERSION >= 7050
-
namespace tensorflow {
namespace detail {
}
}
+// Helper to set all tensor entries to a specific value.
+template <typename T>
+__global__ void SetToValue(const int count, T* ptr, T value) {
+ // Check that the grid is one dimensional and index doesn't overflow.
+ assert(blockDim.y == 1 && blockDim.z == 1);
+ assert(blockDim.x * gridDim.x / blockDim.x == gridDim.x);
+ for (int i : CudaGridRangeX(count)) {
+ ptr[i] = value;
+ }
+}
+
namespace detail {
// Helper function for atomic accumulation implemented as CAS.
template <typename T, typename F>
}));
}
+// Overload of above function for half. Note that we don't have
+// atomicCAS() for anything less than 32 bits, so we need to include the
+// other 16 bits in the operation.
+//
+// This version is going to be very slow
+// under high concurrency, since most threads will be spinning on failing
+// their compare-and-swap tests. (The fact that we get false sharing on the
+// neighboring fp16 makes this even worse.) If you are doing a large reduction,
+// you are much better off with doing the intermediate steps in fp32 and then
+// switching to fp16 as late as you can in the calculations.
+//
+// Note: Assumes little endian.
+template <typename F>
+__device__ Eigen::half CudaAtomicCasHelper(Eigen::half* ptr, F accumulate) {
+#if defined(__BYTE_ORDER__) && defined(__ORDER_LITTLE_ENDIAN__)
+ static_assert(__BYTE_ORDER__ == __ORDER_LITTLE_ENDIAN__, "Not little endian");
+#endif
+ namespace half_impl = Eigen::half_impl;
+ intptr_t intptr = reinterpret_cast<intptr_t>(ptr);
+ assert(!(intptr & 0x1)); // should be 2-aligned.
+ if (intptr & 0x2) {
+ // The half is in the second part of the uint32 (upper 16 bits).
+ uint32* address = reinterpret_cast<uint32*>(intptr - 2);
+ uint32 result = CudaAtomicCasHelper(address, [accumulate](uint32 arg) {
+ unsigned short high = static_cast<unsigned short>(arg >> 16);
+ Eigen::half acc = accumulate(half_impl::raw_uint16_to_half(high));
+ return (static_cast<uint32>(acc.x) << 16) | (arg & 0xffff);
+ });
+ return half_impl::raw_uint16_to_half(static_cast<uint16>(result >> 16));
+ } else {
+ // The half is in the first part of the uint32 (lower 16 bits).
+ uint32* address = reinterpret_cast<uint32*>(intptr);
+ uint32 result = CudaAtomicCasHelper(address, [accumulate](uint32 arg) {
+ unsigned short low = static_cast<unsigned short>(arg & 0xffff);
+ Eigen::half acc = accumulate(half_impl::raw_uint16_to_half(low));
+ return (arg & 0xffff0000) | static_cast<uint32>(acc.x);
+ });
+ return half_impl::raw_uint16_to_half(static_cast<uint16>(result & 0xffff));
+ }
+}
+
template <typename From, typename To>
using ToTypeIfConvertible =
typename std::enable_if<std::is_convertible<From, To>::value, To>::type;
__device__ detail::ToTypeIfConvertible<U, T> CudaAtomicAdd(T* ptr, U value) {
return atomicAdd(ptr, value);
}
+
+__device__ inline Eigen::half CudaAtomicAdd(Eigen::half* ptr,
+ Eigen::half value) {
+ return detail::CudaAtomicCasHelper(
+ ptr, [value](Eigen::half a) { return a + value; });
+}
+
+
#if __CUDA_ARCH__ < 600
__device__ inline double CudaAtomicAdd(double* ptr, double value) {
return detail::CudaAtomicCasHelper(ptr,
return result;
}
#endif
-
+// CudaAtomicAdd
+// Specializations of CudaAtomicAdd for complex types, which CudaAtomicAdd does
+// not support. We treat a std::complex<T>* as a T* (the C++ standard section
+// 26.4.4 allows this explicitly) and atomic add the real and imaginary
+// components individually. The operation as a whole is not atomic, but we can
+// safely treat the components independently for the purpose of accumulating.
+__device__ inline std::complex<float> CudaAtomicAdd(std::complex<float>* ptr,
+ std::complex<float> value) {
+ auto ptr_scalar = reinterpret_cast<float*>(ptr);
+ return std::complex<float>(CudaAtomicAdd(ptr_scalar, value.real()),
+ CudaAtomicAdd(ptr_scalar + 1, value.imag()));
+}
+
+__device__ inline std::complex<double> CudaAtomicAdd(
+ std::complex<double>* ptr, std::complex<double> value) {
+ auto ptr_scalar = reinterpret_cast<double*>(ptr);
+ return std::complex<double>(CudaAtomicAdd(ptr_scalar, value.real()),
+ CudaAtomicAdd(ptr_scalar + 1, value.imag()));
+}
+
+// CudaAtomicSub
template <typename T, typename U>
__device__ detail::ToTypeIfConvertible<U, T> CudaAtomicSub(T* ptr, U value) {
return atomicSub(ptr, value);
}
+
// Specializations of substraction which add the negative value.
__device__ inline float CudaAtomicSub(float* ptr, float value) {
return CudaAtomicAdd(ptr, -value);
}
+
__device__ inline double CudaAtomicSub(double* ptr, double value) {
return CudaAtomicAdd(ptr, -value);
}
+
__device__ inline tensorflow::uint64 CudaAtomicSub(tensorflow::uint64* ptr,
tensorflow::uint64 value) {
return CudaAtomicAdd(ptr, -value);
}
+__device__ inline Eigen::half CudaAtomicSub(Eigen::half* ptr,
+ Eigen::half value) {
+ return detail::CudaAtomicCasHelper(
+ ptr, [value](Eigen::half a) { return a - value; });
+}
+
+// CudaAtomicMax
template <typename T, typename U>
__device__ detail::ToTypeIfConvertible<U, T> CudaAtomicMax(T* ptr, U value) {
return atomicMax(ptr, value);
}
+
+__device__ inline float CudaAtomicMax(float* ptr, float value) {
+ return detail::CudaAtomicCasHelper(
+ ptr, [value](float a) { return max(a, value); });
+}
+
+__device__ inline double CudaAtomicMax(double* ptr, double value) {
+ return detail::CudaAtomicCasHelper(
+ ptr, [value](double a) { return max(a, value); });
+}
+
+__device__ inline Eigen::half CudaAtomicMax(Eigen::half* ptr,
+ Eigen::half value) {
+ return detail::CudaAtomicCasHelper(
+ ptr, [value](Eigen::half a) { return max(a, value); });
+}
+
#if __CUDA_ARCH__ < 320
__device__ inline tensorflow::uint64 CudaAtomicMax(tensorflow::uint64* ptr,
tensorflow::uint64 value) {
}
#endif
+// CudaAtomicMin
+template <typename T, typename U>
+__device__ detail::ToTypeIfConvertible<U, T> CudaAtomicMin(T* ptr, U value) {
+ return atomicMin(ptr, value);
+}
+
+__device__ inline float CudaAtomicMin(float* ptr, float value) {
+ return detail::CudaAtomicCasHelper(
+ ptr, [value](float a) { return min(a, value); });
+}
+
+__device__ inline double CudaAtomicMin(double* ptr, double value) {
+ return detail::CudaAtomicCasHelper(
+ ptr, [value](double a) { return min(a, value); });
+}
+
+__device__ inline Eigen::half CudaAtomicMin(Eigen::half* ptr,
+ Eigen::half value) {
+ return detail::CudaAtomicCasHelper(
+ ptr, [value](Eigen::half a) { return min(a, value); });
+}
+
+#if __CUDA_ARCH__ < 320
+__device__ inline tensorflow::uint64 CudaAtomicMin(tensorflow::uint64* ptr,
+ tensorflow::uint64 value) {
+ return detail::CudaAtomicCasHelper(
+ ptr, [value](tensorflow::uint64 a) { return min(a, value); });
+}
+#endif
+
+// CudaAtomicMul
template <typename T, typename U>
__device__ detail::ToTypeIfConvertible<U, T> CudaAtomicMul(T* ptr, U value) {
return detail::CudaAtomicCasHelper(ptr, [value](T a) { return a * value; });
}
+
+// CudaAtomicDiv
template <typename T, typename U>
__device__ detail::ToTypeIfConvertible<U, T> CudaAtomicDiv(T* ptr, U value) {
return detail::CudaAtomicCasHelper(ptr, [value](T a) { return a / value; });
CudaShuffleXorSync(mask, static_cast<uint16>(value), lane_mask, width));
}
-namespace detail {
-// Overload of above function for half. Note that we don't have
-// atomicCAS() for anything less than 32 bits, so we need to include the
-// other 16 bits in the operation.
-//
-// This version is going to be very slow
-// under high concurrency, since most threads will be spinning on failing
-// their compare-and-swap tests. (The fact that we get false sharing on the
-// neighboring fp16 makes this even worse.) If you are doing a large reduction,
-// you are much better off with doing the intermediate steps in fp32 and then
-// switching to fp16 as late as you can in the calculations.
-//
-// Note: Assumes little endian.
-template <typename F>
-__device__ Eigen::half CudaAtomicCasHelper(Eigen::half* ptr, F accumulate) {
-#if defined(__BYTE_ORDER__) && defined(__ORDER_LITTLE_ENDIAN__)
- static_assert(__BYTE_ORDER__ == __ORDER_LITTLE_ENDIAN__, "Not little endian");
-#endif
- namespace half_impl = Eigen::half_impl;
- intptr_t intptr = reinterpret_cast<intptr_t>(ptr);
- assert(!(intptr & 0x1)); // should be 2-aligned.
- if (intptr & 0x2) {
- // The half is in the second part of the uint32 (upper 16 bits).
- uint32* address = reinterpret_cast<uint32*>(intptr - 2);
- uint32 result = CudaAtomicCasHelper(address, [accumulate](uint32 arg) {
- unsigned short high = static_cast<unsigned short>(arg >> 16);
- Eigen::half acc = accumulate(half_impl::raw_uint16_to_half(high));
- return (static_cast<uint32>(acc.x) << 16) | (arg & 0xffff);
- });
- return half_impl::raw_uint16_to_half(static_cast<uint16>(result >> 16));
- } else {
- // The half is in the first part of the uint32 (lower 16 bits).
- uint32* address = reinterpret_cast<uint32*>(intptr);
- uint32 result = CudaAtomicCasHelper(address, [accumulate](uint32 arg) {
- unsigned short low = static_cast<unsigned short>(arg & 0xffff);
- Eigen::half acc = accumulate(half_impl::raw_uint16_to_half(low));
- return (arg & 0xffff0000) | static_cast<uint32>(acc.x);
- });
- return half_impl::raw_uint16_to_half(static_cast<uint16>(result & 0xffff));
- }
-}
-} // namespace detail
-
-__device__ inline Eigen::half CudaAtomicAdd(Eigen::half* ptr,
- Eigen::half value) {
- return detail::CudaAtomicCasHelper(
- ptr, [value](Eigen::half a) { return a + value; });
-}
-__device__ inline Eigen::half CudaAtomicSub(Eigen::half* ptr,
- Eigen::half value) {
- return detail::CudaAtomicCasHelper(
- ptr, [value](Eigen::half a) { return a - value; });
-}
-
namespace cuda_helper {
template <typename IntType>
__device__ IntType upper_bound(IntType* first, IntType count, IntType val) {
The first time you call an Estimator's `train` method, TensorFlow saves a
checkpoint to the `model_dir`. Each subsequent call to the Estimator's
-`train`, `eval`, or `predict` method causes the following:
+`train`, `evaluate`, or `predict` method causes the following:
1. The Estimator builds the model's
[graph](https://developers.google.com/machine-learning/glossary/#graph)
To run experiments in which you train and compare slightly different
versions of a model, save a copy of the code that created each
-`model-dir`, possibly by creating a separate git branch for each version.
+`model_dir`, possibly by creating a separate git branch for each version.
This separation will keep your checkpoints recoverable.
## Summary
```
* The `units` parameter defines the number of output neurons in a given layer.
-* The `activation` parameter defines the [activation function](https://developers.google.com/machine-learning/glossary/#a) —
+* The `activation` parameter defines the [activation function](https://developers.google.com/machine-learning/glossary/#activation_function) —
[Relu](https://developers.google.com/machine-learning/glossary/#ReLU) in this
case.
dimension in `operand`.
The algorithm goes as follows for each batch in `operand` \\(x\\) that
-contains `m` elements with `w` and `h` as the size of spatial dimensions (
-assuming `operand` is an 4 dimensional array):
+contains `m` elements with `w` and `h` as the size of spatial dimensions
+(assuming `operand` is an 4 dimensional array):
- Calculates batch mean \\(\mu_l\\) for each feature `l` in feature dimension:
\\(\mu_l=\frac{1}{mwh}\sum_{i=1}^m\sum_{j=1}^w\sum_{k=1}^h x_{ijkl}\\)
operation from a data shape to a target shape. The dimensions must match, and
the conversion is an element-wise one; e.g. `s32` elements become `f32` elements
via bitcast routine. Bitcast is implemented as a low-level cast, so machines
-with different floating point representations will give different results.
+with different floating-point representations will give different results.
<b> `BitcastConvertType(operand, new_element_type)` </b>
: : : concatenated between the `operands`. :
With the exception of `dimension` all dimensions must be the same. This is
-because XLA does not support "ragged" arrays Also note that rank-0 values
+because XLA does not support "ragged" arrays. Also note that rank-0 values
cannot be concatenated (as it's impossible to name the dimension along which the
concatenation occurs).
window that moves across the base area.
The `window_strides` argument specifies the stride of the convolutional window
-in the spatial dimensions. For example, if the stride in a the first spatial
+in the spatial dimensions. For example, if the stride in the first spatial
dimension is 3, then the window can only be placed at coordinates where the
first spatial index is divisible by 3.
operand. `broadcast_dimensions` maps the dimensions of the lower-rank shape to
the dimensions of the higher-rank shape. The unmapped dimensions of the expanded
shape are filled with dimensions of size one. Degenerate-dimension broadcasting
-then broadcasts the shapes along these degenerate dimension to equalize the
+then broadcasts the shapes along these degenerate dimensions to equalize the
shapes of both operands. The semantics are described in detail on the
@{$broadcasting$broadcasting page}.
```
Nested tuple shapes are not supported. For an empty tuple shape, the Infeed
-operation is effectively a nop and proceeds without reading any data from the
+operation is effectively a no-op and proceeds without reading any data from the
Infeed of the device.
> Note: We plan to allow multiple Infeed operations without a total order, in
`PaddingConfig` is a repeated field of `PaddingConfigDimension`, which contains
three fields for each dimension: `edge_padding_low`, `edge_padding_high`, and
-`interior_padding`. `edge_padding_low` and `edge_padding_high` specifies the
+`interior_padding`. `edge_padding_low` and `edge_padding_high` specify the
amount of padding added at the low-end (next to index 0) and the high-end (next
to the highest index) of each dimension respectively. The amount of edge padding
can be negative -- the absolute value of negative padding indicates the number
padding occurs logically before edge padding, so in the case of negative edge
padding elements are removed from the interior-padded operand. This operation is
a no-op if the edge padding pairs are all (0, 0) and the interior padding values
-are all 0. Figure below shows examples of different `edge_padding` and
-`interior_padding` values for a two dimensional array.
+are all 0. The figure below shows examples of different `edge_padding` and
+`interior_padding` values for a two-dimensional array.
<div style="width:95%; margin:auto; margin-bottom:10px; margin-top:20px;">
<img style="width:100%" src="https://www.tensorflow.org/images/ops_pad.png">
$ saved_model_cli show --dir \
/tmp/saved_model_dir --tag_set serve --signature_def serving_default
The given SavedModel SignatureDef contains the following input(s):
-inputs['x'] tensor_info:
- dtype: DT_FLOAT
- shape: (-1, 1)
- name: x:0
+ inputs['x'] tensor_info:
+ dtype: DT_FLOAT
+ shape: (-1, 1)
+ name: x:0
The given SavedModel SignatureDef contains the following output(s):
-outputs['y'] tensor_info:
- dtype: DT_FLOAT
- shape: (-1, 1)
- name: y:0
+ outputs['y'] tensor_info:
+ dtype: DT_FLOAT
+ shape: (-1, 1)
+ name: y:0
Method name is: tensorflow/serving/predict
```
MetaGraphDef with tag-set: 'serve' contains the following SignatureDefs:
signature_def['classify_x2_to_y3']:
-The given SavedModel SignatureDef contains the following input(s):
-inputs['inputs'] tensor_info:
- dtype: DT_FLOAT
- shape: (-1, 1)
- name: x2:0
-The given SavedModel SignatureDef contains the following output(s):
-outputs['scores'] tensor_info:
- dtype: DT_FLOAT
- shape: (-1, 1)
- name: y3:0
-Method name is: tensorflow/serving/classify
+ The given SavedModel SignatureDef contains the following input(s):
+ inputs['inputs'] tensor_info:
+ dtype: DT_FLOAT
+ shape: (-1, 1)
+ name: x2:0
+ The given SavedModel SignatureDef contains the following output(s):
+ outputs['scores'] tensor_info:
+ dtype: DT_FLOAT
+ shape: (-1, 1)
+ name: y3:0
+ Method name is: tensorflow/serving/classify
...
signature_def['serving_default']:
-The given SavedModel SignatureDef contains the following input(s):
-inputs['x'] tensor_info:
- dtype: DT_FLOAT
- shape: (-1, 1)
- name: x:0
-The given SavedModel SignatureDef contains the following output(s):
-outputs['y'] tensor_info:
- dtype: DT_FLOAT
- shape: (-1, 1)
- name: y:0
-Method name is: tensorflow/serving/predict
+ The given SavedModel SignatureDef contains the following input(s):
+ inputs['x'] tensor_info:
+ dtype: DT_FLOAT
+ shape: (-1, 1)
+ name: x:0
+ The given SavedModel SignatureDef contains the following output(s):
+ outputs['y'] tensor_info:
+ dtype: DT_FLOAT
+ shape: (-1, 1)
+ name: y:0
+ Method name is: tensorflow/serving/predict
```
of tensors or other objects, such as `tf.Variable` instances.
By default every `tf.Variable` gets placed in the following two collections:
+
* `tf.GraphKeys.GLOBAL_VARIABLES` --- variables that can be shared across
-multiple devices,
- * `tf.GraphKeys.TRAINABLE_VARIABLES`--- variables for which TensorFlow will
+ multiple devices,
+ * `tf.GraphKeys.TRAINABLE_VARIABLES` --- variables for which TensorFlow will
calculate gradients.
If you don't want a variable to be trainable, add it to the
'--window_size_ms',
type=float,
default=30.0,
- help='How long each spectrogram timeslice is.',
- )
+ help='How long each spectrogram timeslice is.',)
parser.add_argument(
'--window_stride_ms',
type=float,
default=10.0,
- help='How far to move in time between spectogram timeslices.',
- )
+ help='How far to move in time between spectogram timeslices.',)
parser.add_argument(
'--dct_coefficient_count',
type=int,
"""
device1 = pydev.canonical_name(device1)
device2 = pydev.canonical_name(device2)
- self.assertEqual(device1, device2, "Devices %s and %s are not equal. %s" %
+ self.assertEqual(device1, device2,
+ "Devices %s and %s are not equal. %s" %
(device1, device2, msg))
# Fix Python 3 compatibility issues
units = 2
model = keras.models.Sequential()
- inputs = keras.layers.Dense(
- embedding_dim, input_shape=(timesteps, embedding_dim))
+ inputs = keras.layers.Dense(embedding_dim,
+ input_shape=(timesteps, embedding_dim))
model.add(inputs)
layer = keras.layers.LSTM(units, return_sequences=True)
model.add(layer)
with self.test_session() as sess:
x = random_ops.random_normal(shape=(2, 2, 3, 4))
- # This LinearOperatorDiag will be brodacast to (2, 2, 3, 3) during solve
+ # This LinearOperatorDiag will be broadcast to (2, 2, 3, 3) during solve
# and matmul with 'x' as the argument.
diag = random_ops.random_uniform(shape=(2, 1, 3))
operator = linalg.LinearOperatorDiag(diag, is_self_adjoint=True)
return constant_op.constant(
np_values, shape=input_shape, dtype=dtype), np_values
- def _segmentReduce(self, indices, x, op1, op2=None, num_segments=None):
+ def _segmentReduce(self, indices, x, op1, op2=None, num_segments=None,
+ initial_value=0):
if not x.size:
return np.array([])
indices = np.asarray(indices)
else:
output[index] = x_flat[i]
# zero initialize values that are still uncalcuated.
- # output = [o if o is not None else np.zeros(slice_shape) for o in output]
- if not op1 == np.max:
- output = [o if o is not None else np.zeros(slice_shape) for o in output]
- else:
- zeroslice = np.zeros(slice_shape)
- zeroslice.fill(dtype.min)
- output = [o if o is not None else zeroslice for o in output]
+ initial_value_slice = np.ones(slice_shape) * initial_value
+ output = [o if o is not None else initial_value_slice for o in output]
if op2 is not None:
output = [op2(o) for o in output]
output = [o.reshape(slice_shape) for o in output]
def _mean_reduce_op(self, x):
return x[0] / x[1] if isinstance(x, tuple) else x
+ def _sqrt_n_reduce_op(self, x):
+ return x[0] / np.sqrt(x[1]) if isinstance(x, tuple) else x
+
class SegmentReductionOpTest(SegmentReductionHelper):
self.assertAllClose(jacob_t, jacob_n)
-class UnsortedSegmentSumTest(SegmentReductionHelper):
+class UnsortedSegmentTest(SegmentReductionHelper):
+
+ def __init__(self, methodName='runTest'):
+ # Each item is np_op1, np_op2, tf_op, initial_value functor
+ self.ops_list = [(np.add, None,
+ math_ops.unsorted_segment_sum, lambda t: 0),
+ (self._mean_cum_op, self._mean_reduce_op,
+ math_ops.unsorted_segment_mean, lambda t: 0),
+ (self._mean_cum_op, self._sqrt_n_reduce_op,
+ math_ops.unsorted_segment_sqrt_n, lambda t: 0),
+ (np.ndarray.__mul__, None,
+ math_ops.unsorted_segment_prod, lambda t: 1),
+ (np.minimum, None,
+ math_ops.unsorted_segment_min, lambda t: t.max),
+ (np.maximum, None,
+ math_ops.unsorted_segment_max, lambda t: t.min)]
+
+ # A subset of ops has been enabled for complex numbers
+ self.complex_ops_list = [(np.add, None,
+ math_ops.unsorted_segment_sum, lambda t: 0)]
+ self.differentiable_dtypes = [dtypes_lib.float16, dtypes_lib.float32,
+ dtypes_lib.float64]
+ self.all_dtypes = (self.differentiable_dtypes +
+ [dtypes_lib.bfloat16,
+ dtypes_lib.int64, dtypes_lib.int32,
+ dtypes_lib.complex64, dtypes_lib.complex128])
+ super(UnsortedSegmentTest, self).__init__(methodName=methodName)
def testValues(self):
- dtypes = [
- dtypes_lib.float32, dtypes_lib.float64, dtypes_lib.int64,
- dtypes_lib.int32, dtypes_lib.complex64, dtypes_lib.complex128
- ]
indices_flat = np.array([0, 4, 0, 8, 3, 8, 4, 7, 7, 3])
num_segments = 12
for indices in indices_flat, indices_flat.reshape(5, 2):
shape = indices.shape + (2,)
- for dtype in dtypes:
- with self.test_session(use_gpu=True):
- tf_x, np_x = self._input(shape, dtype=dtype)
- np_ans = self._segmentReduce(
- indices, np_x, np.add, op2=None, num_segments=num_segments)
- s = math_ops.unsorted_segment_sum(
- data=tf_x, segment_ids=indices, num_segments=num_segments)
- tf_ans = s.eval()
- self.assertAllClose(np_ans, tf_ans)
- self.assertShapeEqual(np_ans, s)
+ for dtype in self.all_dtypes:
+ ops_list = self.complex_ops_list if dtype.is_complex else self.ops_list
+ tf_x, np_x = self._input(shape, dtype=dtype)
+ for use_gpu in [True, False]:
+ with self.test_session(use_gpu=True):
+ for np_op1, np_op2, tf_op, init_op in ops_list:
+ # sqrt_n doesn't support integers
+ if (np_op2 == self._sqrt_n_reduce_op and dtype.is_integer):
+ continue
+ # todo(philjd): enable this test once real_div supports bfloat16
+ if (np_op2 in [self._sqrt_n_reduce_op, self._mean_reduce_op] and
+ dtype == dtypes_lib.bfloat16):
+ continue
+ np_ans = self._segmentReduce(
+ indices, np_x, np_op1, np_op2, num_segments=num_segments,
+ initial_value=init_op(dtype))
+ s = tf_op(tf_x, segment_ids=indices, num_segments=num_segments)
+ tf_ans = s.eval()
+ if dtype is dtypes_lib.bfloat16:
+ tf_ans = tf_ans.astype(np.float32)
+ self.assertAllClose(np_ans, tf_ans)
+ self.assertShapeEqual(np_ans, s)
def testNumSegmentsTypes(self):
dtypes = [dtypes_lib.int32, dtypes_lib.int64]
self.assertAllClose(np_ans, tf_ans)
self.assertShapeEqual(np_ans, s)
- def testGradientSegmentSum(self):
+ def testGradients(self):
num_cols = 2
- indices_flat = np.array([0, 4, 0, 8, 3, 8, 4, 7, 7, 3])
+ indices_flat = np.array([0, 4, 0, -1, 3, -1, 4, 7, 7, 3])
num_segments = max(indices_flat) + 3
- for dtype in [dtypes_lib.float32, dtypes_lib.float64, dtypes_lib.complex64,
- dtypes_lib.complex128]:
+ for dtype in self.differentiable_dtypes:
+ ops_list = self.complex_ops_list if dtype.is_complex else self.ops_list
for indices in indices_flat, indices_flat.reshape(5, 2):
shape = indices.shape + (num_cols,)
- with self.test_session(use_gpu=True):
- tf_x, np_x = self._input(shape, dtype=dtype)
- s = math_ops.unsorted_segment_sum(
- data=tf_x, segment_ids=indices, num_segments=num_segments)
+ # test CPU and GPU as tf.gather behaves differently on each device
+ for use_gpu in [False, True]:
+ with self.test_session(use_gpu=use_gpu):
+ for _, _, tf_op, _ in ops_list:
+ tf_x, np_x = self._input(shape, dtype=dtype)
+ s = tf_op(tf_x, indices, num_segments)
+ jacob_t, jacob_n = gradient_checker.compute_gradient(
+ tf_x,
+ shape,
+ s, [num_segments, num_cols],
+ x_init_value=np_x,
+ delta=1)
+ self.assertAllClose(jacob_t, jacob_n)
+
+ def testProdGrad(self):
+ # additional test for the prod gradient to ensure correct handling of zeros
+ values = np.array([0, 0, 1, 0, 2, 2, 3, 3, 3], dtype=np.float32)
+ indices = np.array([0, 0, 0, 1, 1, 1, 2, 2, 2], dtype=np.int32)
+ indices_neg = np.array([-1, 0, 0, -1, 1, 1, -1, 2, 2], dtype=np.int32)
+ values_tf = constant_op.constant(values)
+ # ground truth partial derivatives
+ gradients_indices = np.zeros((9, 3), dtype=np.float32)
+ gradients_indices_neg = np.zeros((9, 3), dtype=np.float32)
+ # the derivative w.r.t. to the other segments is zero, so here we only
+ # explicitly set the grad values for the corresponding segment
+ gradients_indices[range(9), indices] = [0, 0, 0, 4, 0, 0, 9, 9, 9]
+ gradients_indices_neg[range(9), indices_neg] = [0, 1, 0, 0, 2, 2, 0, 3, 3]
+ for use_gpu in [False, True]:
+ with self.test_session(use_gpu=use_gpu):
+ for ind, grad_gt in [(indices, gradients_indices),
+ (indices_neg, gradients_indices_neg)]:
+ s = math_ops.unsorted_segment_prod(values_tf,
+ constant_op.constant(ind), 3)
jacob_t, jacob_n = gradient_checker.compute_gradient(
- tf_x,
- shape,
- s, [num_segments, num_cols],
- x_init_value=np_x,
- delta=1)
- self.assertAllClose(jacob_t, jacob_n)
+ values_tf, (9,), s, (3,), x_init_value=values, delta=1)
+ self.assertAllClose(jacob_t, jacob_n)
+ self.assertAllClose(jacob_t, grad_gt)
def testGradientMatchesSegmentSum(self):
# Strategy: compute the gradient for UnsortedSegmentSum and SegmentSum
num_cols = 2
shape = [n, num_cols]
num_segments = max(indices) + 1
- for dtype in [dtypes_lib.float32, dtypes_lib.float64, dtypes_lib.complex64,
- dtypes_lib.complex128]:
+ for dtype in self.differentiable_dtypes:
with self.test_session(use_gpu=True):
tf_x, np_x = self._input(shape, dtype=dtype)
# Results from UnsortedSegmentSum
unsorted.eval()
def testEmptySecondDimension(self):
- dtypes = [
- np.float32, np.float64, np.int64, np.int32, np.complex64, np.complex128
- ]
+ dtypes = [np.float16, np.float32, np.float64, np.int64, np.int32,
+ np.complex64, np.complex128]
with self.test_session(use_gpu=True):
for dtype in dtypes:
for itype in (np.int32, np.int64):
unsorted = math_ops.unsorted_segment_sum(data, segment_ids, 2)
self.assertAllEqual(unsorted.eval(), np.zeros((2, 0), dtype=dtype))
- def testGradientSegmentMax(self):
- num_cols = 2
- indices_flat = np.array([0, 4, 0, 8, 3, 8, 4, 7, 7, 3])
- num_segments = max(indices_flat) + 3
- for indices in indices_flat, indices_flat.reshape(5, 2):
- shape = indices.shape + (num_cols,)
- with self.test_session(use_gpu=True):
- tf_x, np_x = self._input(shape, dtype=dtypes_lib.float64)
- s = math_ops.unsorted_segment_max(
- data=tf_x, segment_ids=indices, num_segments=num_segments)
- jacob_t, jacob_n = gradient_checker.compute_gradient(
- tf_x,
- shape,
- s,
- [num_segments, num_cols],
- x_init_value=np_x.astype(np.double), delta=1)
- self.assertAllClose(jacob_t, jacob_n)
-
def testDropNegatives(self):
# Note: the test is done by replacing segment_ids with 8 to -1
# for index and replace values generated by numpy with 0.
- dtypes = [
- dtypes_lib.float32, dtypes_lib.float64, dtypes_lib.int64,
- dtypes_lib.int32, dtypes_lib.complex64, dtypes_lib.complex128
- ]
indices_flat = np.array([0, 4, 0, 8, 3, 8, 4, 7, 7, 3])
num_segments = 12
for indices in indices_flat, indices_flat.reshape(5, 2):
shape = indices.shape + (2,)
- for dtype in dtypes:
+ for dtype in self.all_dtypes:
with self.test_session(use_gpu=True):
tf_x, np_x = self._input(shape, dtype=dtype)
np_ans = self._segmentReduce(
v = [1 if x[i] == value.decode('ascii') else 0 for i in range(7000)]
self.assertEqual(count, sum(v))
+ def testInt32Axis(self):
+ for dtype in [np.int32, np.int64]:
+ x = np.array([[1, 0, 0], [1, 0, 0], [2, 0, 0]])
+ with self.test_session() as sess:
+ y0, idx0, count0 = gen_array_ops.unique_with_counts_v2(
+ x, axis=np.array([0], dtype))
+ tf_y0, tf_idx0, tf_count0 = sess.run([y0, idx0, count0])
+ y1, idx1, count1 = gen_array_ops.unique_with_counts_v2(
+ x, axis=np.array([1], dtype))
+ tf_y1, tf_idx1, tf_count1 = sess.run([y1, idx1, count1])
+ self.assertAllEqual(tf_y0, np.array([[1, 0, 0], [2, 0, 0]]))
+ self.assertAllEqual(tf_idx0, np.array([0, 0, 1]))
+ self.assertAllEqual(tf_count0, np.array([2, 1]))
+ self.assertAllEqual(tf_y1, np.array([[1, 0], [1, 0], [2, 0]]))
+ self.assertAllEqual(tf_idx1, np.array([0, 1, 1]))
+ self.assertAllEqual(tf_count1, np.array([1, 2]))
+
+ def testInt32V2(self):
+ # This test is only temporary, once V2 is used
+ # by default, the axis will be wrapped to allow `axis=None`.
+ x = np.random.randint(2, high=10, size=7000)
+ with self.test_session() as sess:
+ y, idx, count = gen_array_ops.unique_with_counts_v2(
+ x, axis=np.array([], np.int32))
+ tf_y, tf_idx, tf_count = sess.run([y, idx, count])
+
+ self.assertEqual(len(x), len(tf_idx))
+ self.assertEqual(len(tf_y), len(np.unique(x)))
+ for i in range(len(x)):
+ self.assertEqual(x[i], tf_y[tf_idx[i]])
+ for value, count in zip(tf_y, tf_count):
+ self.assertEqual(count, np.sum(x == value))
+
if __name__ == '__main__':
test.main()
unique.__doc__ = gen_array_ops.unique.__doc__
+@tf_export("unique_with_counts")
+def unique_with_counts(x, out_idx=dtypes.int32, name=None):
+ # TODO(yongtang): switch to v2 once API deprecation
+ # period (3 weeks) pass.
+ # TODO(yongtang): The documentation should also
+ # be updated when switch to v2.
+ return gen_array_ops.unique_with_counts(x, out_idx, name)
+
+
+unique_with_counts.__doc__ = gen_array_ops.unique_with_counts.__doc__
+
+
@tf_export("split")
def split(value, num_or_size_splits, axis=0, num=None, name="split"):
"""Splits a tensor into sub tensors.
self.assertAllEqual(truth, popcnt_result)
def testInvertOp(self):
- dtype_list = [
- dtypes.int8, dtypes.int16, dtypes.int32, dtypes.int64, dtypes.uint8,
- dtypes.uint16, dtypes.uint32, dtypes.uint64
- ]
+ dtype_list = [dtypes.int8, dtypes.int16, dtypes.int32, dtypes.int64,
+ dtypes.uint8, dtypes.uint16, dtypes.uint32, dtypes.uint64]
inputs = [0, 5, 3, 14]
with self.test_session(use_gpu=True) as sess:
for dtype in dtype_list:
@compatibility{eager} returns None
Raises:
- InvalidArgumentError if the check can be performed immediately and
- `x == y` is False. The check can be performed immediately during
- eager execution or if `x` and `y` are statically known.
+ InvalidArgumentError: if the check can be performed immediately and
+ `x == y` is False. The check can be performed immediately during eager
+ execution or if `x` and `y` are statically known.
"""
message = message or ''
with ops.name_scope(name, 'assert_equal', [x, y, data]):
name=None, weights=None):
"""Computes the confusion matrix from predictions and labels.
- Calculate the Confusion Matrix for a pair of prediction and
- label 1-D int arrays.
-
The matrix columns represent the prediction labels and the rows represent the
real labels. The confusion matrix is always a 2-D array of shape `[n, n]`,
where `n` is the number of valid labels for a given classification task. Both
prediction and labels must be 1-D arrays of the same shape in order for this
function to work.
- If `num_classes` is None, then `num_classes` will be set to the one plus
- the maximum value in either predictions or labels.
- Class labels are expected to start at 0. E.g., if `num_classes` was
- three, then the possible labels would be `[0, 1, 2]`.
+ If `num_classes` is `None`, then `num_classes` will be set to one plus the
+ maximum value in either predictions or labels. Class labels are expected to
+ start at 0. For example, if `num_classes` is 3, then the possible labels
+ would be `[0, 1, 2]`.
If `weights` is not `None`, then each prediction contributes its
corresponding weight to the total value of the confusion matrix cell.
weights: An optional `Tensor` whose shape matches `predictions`.
Returns:
- A k X k matrix representing the confusion matrix, where k is the number of
- possible labels in the classification task.
+ A `Tensor` of type `dtype` with shape `[n, n]` representing the confusion
+ matrix, where `n` is the number of possible labels in the classification
+ task.
Raises:
ValueError: If both predictions and labels are not 1-D vectors and have
weights = math_ops.cast(weights, dtype)
shape = array_ops.stack([num_classes, num_classes])
- indices = array_ops.transpose(array_ops.stack([labels, predictions]))
+ indices = array_ops.stack([labels, predictions], axis=1)
values = (array_ops.ones_like(predictions, dtype)
if weights is None else weights)
cm_sparse = sparse_tensor.SparseTensor(
# Compute x for p <= exp(-2): x = z - log(z)/z - (1/z) P(1/z) / Q(1/z),
# where z = sqrt(-2. * log(p)), and P/Q are chosen between two different
- # arrays based on wether p < exp(-32).
+ # arrays based on whether p < exp(-32).
z = math_ops.sqrt(-2. * math_ops.log(sanitized_mcp))
first_term = z - math_ops.log(z) / z
second_term_small_p = (_create_polynomial(1. / z, p2)
ZerosLike # TODO(josh11b): Use this instead of the Python version.
Unique
UniqueV2
+UniqueWithCounts
+UniqueWithCountsV2
Unpack
# candidate_sampling_ops
def _rot180():
return array_ops.reverse_v2(images, [1, 2])
-
def _rot270():
return array_ops.reverse_v2(array_ops.transpose(images, [0, 2, 1, 3]), [2])
result.set_shape([shape[0], None, None, shape[3]])
return result
-
@tf_export('image.transpose_image')
def transpose_image(image):
"""Transpose image(s) by swapping the height and width dimension.
with self.assertRaisesRegexp(ValueError, "must be three-dimensional"):
op(p_wrong_rank)
+
def testRot90GroupOrder(self):
image = np.arange(24, dtype=np.uint8).reshape([2, 4, 3])
with self.test_session(use_gpu=True):
y_np = np.rot90(image, k=k, axes=(1, 2))
self.assertAllEqual(y_np, y_tf.eval({k_placeholder: k}))
-
class RandomFlipTest(test_util.TensorFlowTestCase):
def testRandomLeftRight(self):
operator = LinearOperatorDiag(diag)
# Create a shape [2, 1, 4, 2] vector. Note that this shape is compatible
- # since the batch dimensions, [2, 1], are brodcast to
+ # since the batch dimensions, [2, 1], are broadcast to
# operator.batch_shape = [2, 3].
y = tf.random_normal(shape=[2, 1, 4, 2])
x = operator.solve(y)
Args:
multi_class_labels: `[batch_size, num_classes]` target integer labels in
- `(0, 1)`.
+ `{0, 1}`.
logits: Float `[batch_size, num_classes]` logits outputs of the network.
weights: Optional `Tensor` whose rank is either 0, or the same rank as
`labels`, and must be broadcastable to `labels` (i.e., all dimensions must
dim0), None, None, None)
-def _SegmentMinOrMaxGrad(op, grad, is_sorted):
- """Gradient for SegmentMin and (unsorted) SegmentMax.
-
- They share similar code.
- """
- zeros = array_ops.zeros(
- array_ops.shape(op.inputs[0]), dtype=op.inputs[0].dtype)
-
+def _SegmentMinOrMaxGrad(op, grad):
+ """ Gradient for SegmentMin and SegmentMax. """
+ zeros = array_ops.zeros_like(op.inputs[0], dtype=op.inputs[0].dtype)
# Get the number of selected (minimum or maximum) elements in each segment.
gathered_outputs = array_ops.gather(op.outputs[0], op.inputs[1])
is_selected = math_ops.equal(op.inputs[0], gathered_outputs)
- if is_sorted:
- num_selected = math_ops.segment_sum(
- math_ops.cast(is_selected, grad.dtype), op.inputs[1])
- else:
- num_selected = math_ops.unsorted_segment_sum(
- math_ops.cast(is_selected, grad.dtype), op.inputs[1], op.inputs[2])
-
+ num_selected = math_ops.segment_sum(math_ops.cast(is_selected, grad.dtype),
+ op.inputs[1])
# Compute the gradient for each segment. The gradient for the ith segment is
# divided evenly among the selected elements in that segment.
weighted_grads = math_ops.div(grad, num_selected)
gathered_grads = array_ops.gather(weighted_grads, op.inputs[1])
-
- if is_sorted:
- return array_ops.where(is_selected, gathered_grads, zeros), None
- else:
- return array_ops.where(is_selected, gathered_grads, zeros), None, None
+ return array_ops.where(is_selected, gathered_grads, zeros), None
@ops.RegisterGradient("SegmentMin")
def _SegmentMinGrad(op, grad):
"""Gradient for SegmentMin."""
- return _SegmentMinOrMaxGrad(op, grad, True)
+ return _SegmentMinOrMaxGrad(op, grad)
@ops.RegisterGradient("SegmentMax")
def _SegmentMaxGrad(op, grad):
"""Gradient for SegmentMax."""
- return _SegmentMinOrMaxGrad(op, grad, True)
+ return _SegmentMinOrMaxGrad(op, grad)
+
+
+def _GatherDropNegatives(params, ids, zero_clipped_indices=None,
+ is_positive=None):
+ """ Helper function for unsorted segment ops. Gathers params for
+ positive segment ids and gathers 0 for inputs with negative segment id.
+ Also returns the clipped indices and a boolean mask with the same shape
+ as ids where a positive id is masked as true. With this, the latter two
+ can be passed as arguments to this function to reuse them.
+ """
+ if zero_clipped_indices is None:
+ zero_clipped_indices = math_ops.maximum(ids, array_ops.zeros_like(ids))
+ gathered = array_ops.gather(params, zero_clipped_indices)
+ if is_positive is None:
+ is_positive = math_ops.greater_equal(ids, 0)
+ # tf.where(condition, x, y) requires condition to have the same shape as x
+ # and y.
+ # todo(philjd): remove this if tf.where supports broadcasting (#9284)
+ for _ in range(gathered.shape.ndims - is_positive.shape.ndims):
+ is_positive = array_ops.expand_dims(is_positive, -1)
+ is_positive = (is_positive &
+ array_ops.ones_like(gathered, dtype=dtypes.bool))
+ # replace gathered params of negative indices with 0
+ zero_slice = array_ops.zeros_like(gathered)
+ return (array_ops.where(is_positive, gathered, zero_slice),
+ zero_clipped_indices, is_positive)
+
+
+def _UnsortedSegmentMinOrMaxGrad(op, grad):
+ """ Gradient for UnsortedSegmentMin and UnsortedSegmentMax. """
+ # Get the number of selected (minimum or maximum) elements in each segment.
+ gathered_outputs, zero_clipped_indices, is_positive = \
+ _GatherDropNegatives(op.outputs[0], op.inputs[1])
+ is_selected = math_ops.equal(op.inputs[0], gathered_outputs)
+ is_selected = math_ops.logical_and(is_selected, is_positive)
+ num_selected = math_ops.unsorted_segment_sum(
+ math_ops.cast(is_selected, grad.dtype), op.inputs[1], op.inputs[2])
+ # Compute the gradient for each segment. The gradient for the ith segment is
+ # divided evenly among the selected elements in that segment.
+ weighted_grads = math_ops.div(grad, num_selected)
+ gathered_grads, _, _ = _GatherDropNegatives(weighted_grads, None,
+ zero_clipped_indices,
+ is_positive)
+ zeros = array_ops.zeros_like(gathered_grads)
+ return array_ops.where(is_selected, gathered_grads, zeros), None, None
@ops.RegisterGradient("UnsortedSegmentSum")
def _UnsortedSegmentSumGrad(op, grad):
- """Gradient for SegmentSum."""
- return array_ops.gather(grad, op.inputs[1]), None, None
+ """Gradient for UnsortedSegmentSum."""
+ return _GatherDropNegatives(grad, op.inputs[1])[0], None, None
@ops.RegisterGradient("UnsortedSegmentMax")
def _UnsortedSegmentMaxGrad(op, grad):
- return _SegmentMinOrMaxGrad(op, grad, False)
+ """ Gradient for UnsortedSegmentMax. """
+ return _UnsortedSegmentMinOrMaxGrad(op, grad)
+
+
+@ops.RegisterGradient("UnsortedSegmentMin")
+def _UnsortedSegmentMinGrad(op, grad):
+ """ Gradient for UnsortedSegmentMin. """
+ return _UnsortedSegmentMinOrMaxGrad(op, grad)
+
+
+@ops.RegisterGradient("UnsortedSegmentProd")
+def _UnsortedSegmentProdGrad(op, grad):
+ """ Gradient for UnsortedSegmentProd.
+ The gradient can be expressed for each segment by dividing the segment's
+ product by each element of the segment input tensor, but this approach can't
+ deal with zeros in the input.
+ Unlike reduce_prod we can't use cumsum here as individual segments may have
+ a different number of elements. Therefore we consider three cases:
+ 1) A segment input contains no zeros and we can safely divide by the input
+ tensor.
+ 2) A segment contains exactly one zero. Then the gradient of each input of
+ the segment is zero except for the 0-input, there the gradient is
+ the product of the remaining segment entries.
+ 3) A segment contains at least two zeros. The gradient is zero for all
+ segment inputs.
+ """
+ # Note that unsorted_segment_sum will filter out the negative indices,
+ # so we don't need to do a logical_and with is_positive here
+ is_zero = math_ops.equal(op.inputs[0], 0)
+ num_zeros = gen_math_ops.unsorted_segment_sum(
+ math_ops.cast(is_zero, dtype=dtypes.int32), op.inputs[1], op.inputs[2])
+ # handle case 3 and set the gradient to 0 for segments with more than one
+ # 0 as input
+ grad = array_ops.where(math_ops.greater(num_zeros, 1),
+ array_ops.zeros_like(grad), grad)
+ # replace all zeros with ones and compute the unsorted_segment_prod
+ non_zero_data = array_ops.where(is_zero, array_ops.ones_like(op.inputs[0]),
+ op.inputs[0])
+ non_zero_prod = gen_math_ops.unsorted_segment_prod(
+ non_zero_data, op.inputs[1], op.inputs[2])
+ # clip the indices for gather to be positive
+ zero_clipped_indices = math_ops.maximum(op.inputs[1],
+ array_ops.zeros_like(op.inputs[1]))
+ gathered_prod = array_ops.gather(op.outputs[0], zero_clipped_indices)
+ gathered_non_zero_prod = array_ops.gather(non_zero_prod,
+ zero_clipped_indices)
+ prod_divided_by_el = gathered_prod / op.inputs[0] # May contain nan/inf.
+ # Now fetch the individual results for segments containing 0 and those that
+ # don't. is_zero will also fetch results for entries with negative index
+ # but the following gather_drop_negatives sets the corresponding entry in
+ # grad to 0 for these
+ partial_derivative = array_ops.where(is_zero, gathered_non_zero_prod,
+ prod_divided_by_el)
+ gathered_grad = _GatherDropNegatives(grad, op.inputs[1],
+ zero_clipped_indices)[0]
+ return gathered_grad * partial_derivative, None, None
@ops.RegisterGradient("Abs")
@@segment_mean
@@unsorted_segment_sum
@@unsorted_segment_max
+@@unsorted_segment_min
+@@unsorted_segment_prod
+@@unsorted_segment_sqrt_n
@@sparse_segment_sum
@@sparse_segment_mean
@@sparse_segment_sqrt_n
return cast(x, dtypes.bfloat16, name=name)
+@tf_export("to_complex64")
+def to_complex64(x, name="ToComplex64"):
+ """Casts a tensor to type `complex64`.
+
+ Args:
+ x: A `Tensor` or `SparseTensor`.
+ name: A name for the operation (optional).
+
+ Returns:
+ A `Tensor` or `SparseTensor` with same shape as `x` with type `complex64`.
+
+ Raises:
+ TypeError: If `x` cannot be cast to the `complex64`.
+ """
+ return cast(x, dtypes.complex64, name=name)
+
+
+@tf_export("to_complex128")
+def to_complex128(x, name="ToComplex128"):
+ """Casts a tensor to type `complex128`.
+
+ Args:
+ x: A `Tensor` or `SparseTensor`.
+ name: A name for the operation (optional).
+
+ Returns:
+ A `Tensor` or `SparseTensor` with same shape as `x` with type `complex128`.
+
+ Raises:
+ TypeError: If `x` cannot be cast to the `complex128`.
+ """
+ return cast(x, dtypes.complex128, name=name)
+
+
ops.Tensor._override_operator("__neg__", gen_math_ops.neg)
ops.Tensor._override_operator("__abs__", abs)
# __invert__ corresponds to the ~ operator. Here we follow the numpy convention
]) # [1, 1]
+def _unsorted_segment_N(data, segment_ids, num_segments):
+ """ Helper function for unsorted_segment_mean/_sqrtN. Computes the number
+ of segment entries with 0-entries set to 1 to allow division by N.
+ """
+ # bincount doesn't support negative indices so we use unsorted_segment_sum
+ ones_tensor = array_ops.ones(segment_ids.shape, dtype=data.dtype)
+ N = gen_math_ops.unsorted_segment_sum(ones_tensor, segment_ids, num_segments)
+ # add dimensions for all non-reduced axes
+ ndims_output = data.shape.ndims - segment_ids.shape.ndims
+ broadcast_shape = [num_segments] + [1] * ndims_output
+ N = array_ops.reshape(N, broadcast_shape)
+ return gen_math_ops.maximum(N, 1)
+
+
+@tf_export("unsorted_segment_mean")
+def unsorted_segment_mean(data, segment_ids, num_segments, name=None):
+ r""" Computes the mean along segments of a tensor.
+
+ Read @{$math_ops#segmentation$the section on segmentation} for an explanation
+ of segments.
+
+ This operator is similar to the unsorted segment sum operator found
+ [here](../../../api_docs/python/math_ops.md#UnsortedSegmentSum).
+ Instead of computing the sum over segments, it computes the mean of all
+ entries belonging to a segment such that:
+
+ \\(output_i = 1/N_i \sum data_j\\) where the sum is over `j` such
+ that `segment_ids[j] == i` with \\N_i\\ being the number of occurrences
+ of id \\i\\.
+
+ If there is no entry for a given segment ID `i`, it outputs 0.
+
+ segment_ids: A 1-D tensor whose rank is equal to the rank of `data`'s
+ first dimension.
+
+ output: Has same shape as data, except for dimension 0 which
+ has size `num_segments`.
+ """
+ with ops.name_scope(name, "UnsortedSegmentMean"):
+ data = ops.convert_to_tensor(data)
+ segment_ids = ops.convert_to_tensor(segment_ids)
+ N = _unsorted_segment_N(data, segment_ids, num_segments)
+ summed = gen_math_ops.unsorted_segment_sum(data, segment_ids, num_segments)
+ return summed / N
+
+
+@tf_export("unsorted_segment_sqrt_n")
+def unsorted_segment_sqrt_n(data, segment_ids, num_segments, name=None):
+ r"""Computes the sum along segments of a tensor divided by the sqrt(N).
+
+ Read @{$math_ops#segmentation$the section on segmentation} for an explanation
+ of segments.
+
+ This operator is similar to the unsorted segment sum operator found
+ [here](../../../api_docs/python/math_ops.md#UnsortedSegmentSum).
+ Additionally to computing the sum over segments, it divides the results by
+ sqrt(N).
+
+ \\(output_i = 1/sqrt(N_i) \sum data_j\\) where the sum is over `j` such
+ that `segment_ids[j] == i` with \\N_i\\ being the number of occurrences
+ of id \\i\\.
+
+ If there is no entry for a given segment ID `i`, it outputs 0.
+
+ Note that this op only supports floating point and complex dtypes,
+ due to tf.sqrt only supporting these types.
+
+ segment_ids: A 1-D tensor whose rank is equal to the rank of `data`'s
+ first dimension.
+
+ output: Has same shape as data, except for dimension 0 which
+ has size `num_segments`.
+ """
+ with ops.name_scope(name, "UnsortedSegmentSqrtN"):
+ data = ops.convert_to_tensor(data)
+ segment_ids = ops.convert_to_tensor(segment_ids)
+ N = _unsorted_segment_N(data, segment_ids, num_segments)
+ summed = gen_math_ops.unsorted_segment_sum(data, segment_ids, num_segments)
+ return summed / gen_math_ops.sqrt(N)
+
+
@tf_export("sparse_segment_sum")
def sparse_segment_sum(data, indices, segment_ids, name=None,
num_segments=None):
sampled_losses = nn_ops.softmax_cross_entropy_with_logits(
labels=labels, logits=logits)
# sampled_losses is a [batch_size] tensor.
- return sampled_losses
+ return sampled_losses
\ No newline at end of file
signature_def_key).outputs
-def _show_inputs_outputs(saved_model_dir, tag_set, signature_def_key):
+def _show_inputs_outputs(saved_model_dir, tag_set, signature_def_key, indent=0):
"""Prints input and output TensorInfos.
Prints the details of input and output TensorInfos for the SignatureDef mapped
tag_set: Group of tag(s) of the MetaGraphDef, in string format, separated by
','. For tag-set contains multiple tags, all tags must be passed in.
signature_def_key: A SignatureDef key string.
+ indent: How far (in increments of 2 spaces) to indent each line of output.
"""
meta_graph_def = saved_model_utils.get_meta_graph_def(saved_model_dir,
tag_set)
outputs_tensor_info = _get_outputs_tensor_info_from_meta_graph_def(
meta_graph_def, signature_def_key)
- print('The given SavedModel SignatureDef contains the following input(s):')
+ indent_str = " " * indent
+ def in_print(s):
+ print(indent_str + s)
+
+ in_print('The given SavedModel SignatureDef contains the following input(s):')
for input_key, input_tensor in sorted(inputs_tensor_info.items()):
- print('inputs[\'%s\'] tensor_info:' % input_key)
- _print_tensor_info(input_tensor)
+ in_print(' inputs[\'%s\'] tensor_info:' % input_key)
+ _print_tensor_info(input_tensor, indent+1)
- print('The given SavedModel SignatureDef contains the following output(s):')
+ in_print('The given SavedModel SignatureDef contains the following '
+ 'output(s):')
for output_key, output_tensor in sorted(outputs_tensor_info.items()):
- print('outputs[\'%s\'] tensor_info:' % output_key)
- _print_tensor_info(output_tensor)
+ in_print(' outputs[\'%s\'] tensor_info:' % output_key)
+ _print_tensor_info(output_tensor, indent+1)
- print('Method name is: %s' %
- meta_graph_def.signature_def[signature_def_key].method_name)
+ in_print('Method name is: %s' %
+ meta_graph_def.signature_def[signature_def_key].method_name)
-def _print_tensor_info(tensor_info):
+def _print_tensor_info(tensor_info, indent=0):
"""Prints details of the given tensor_info.
Args:
tensor_info: TensorInfo object to be printed.
+ indent: How far (in increments of 2 spaces) to indent each line output
"""
- print(' dtype: ' +
- {value: key
- for (key, value) in types_pb2.DataType.items()}[tensor_info.dtype])
+ indent_str = " " * indent
+ def in_print(s):
+ print(indent_str + s)
+
+ in_print(' dtype: ' +
+ {value: key
+ for (key, value) in types_pb2.DataType.items()}[tensor_info.dtype])
# Display shape as tuple.
if tensor_info.tensor_shape.unknown_rank:
shape = 'unknown_rank'
dims = [str(dim.size) for dim in tensor_info.tensor_shape.dim]
shape = ', '.join(dims)
shape = '(' + shape + ')'
- print(' shape: ' + shape)
- print(' name: ' + tensor_info.name)
+ in_print(' shape: ' + shape)
+ in_print(' name: ' + tensor_info.name)
def _show_all(saved_model_dir):
signature_def_map = get_signature_def_map(saved_model_dir, tag_set)
for signature_def_key in sorted(signature_def_map.keys()):
print('\nsignature_def[\'' + signature_def_key + '\']:')
- _show_inputs_outputs(saved_model_dir, tag_set, signature_def_key)
+ _show_inputs_outputs(saved_model_dir, tag_set, signature_def_key,
+ indent=1)
def get_meta_graph_def(saved_model_dir, tag_set):
show_msg = (
'Usage examples:\n'
'To show all tag-sets in a SavedModel:\n'
- '$saved_model_cli show --dir /tmp/saved_model\n'
+ '$saved_model_cli show --dir /tmp/saved_model\n\n'
'To show all available SignatureDef keys in a '
'MetaGraphDef specified by its tag-set:\n'
- '$saved_model_cli show --dir /tmp/saved_model --tag_set serve\n'
+ '$saved_model_cli show --dir /tmp/saved_model --tag_set serve\n\n'
'For a MetaGraphDef with multiple tags in the tag-set, all tags must be '
'passed in, separated by \';\':\n'
'$saved_model_cli show --dir /tmp/saved_model --tag_set serve,gpu\n\n'
'To show all inputs and outputs TensorInfo for a specific'
' SignatureDef specified by the SignatureDef key in a'
' MetaGraph.\n'
- '$saved_model_cli show --dir /tmp/saved_model --tag_set serve '
- '--signature_def serving_default\n\n'
- 'To show all available information in the SavedModel\n:'
+ '$saved_model_cli show --dir /tmp/saved_model --tag_set serve'
+ ' --signature_def serving_default\n\n'
+ 'To show all available information in the SavedModel:\n'
'$saved_model_cli show --dir /tmp/saved_model --all')
parser_show = subparsers.add_parser(
'show',
run_msg = ('Usage example:\n'
'To run input tensors from files through a MetaGraphDef and save'
' the output tensors to files:\n'
- '$saved_model_cli show --dir /tmp/saved_model --tag_set serve '
- '--signature_def serving_default '
- '--inputs input1_key=/tmp/124.npz[x],input2_key=/tmp/123.npy '
- '--input_exprs \'input3_key=np.ones(2)\' --input_examples '
- '\'input4_key=[{"id":[26],"weights":[0.5, 0.5]}]\' '
- '--outdir=/out\n\n'
+ '$saved_model_cli show --dir /tmp/saved_model --tag_set serve \\\n'
+ ' --signature_def serving_default \\\n'
+ ' --inputs input1_key=/tmp/124.npz[x],input2_key=/tmp/123.npy '
+ '\\\n'
+ ' --input_exprs \'input3_key=np.ones(2)\' \\\n'
+ ' --input_examples '
+ '\'input4_key=[{"id":[26],"weights":[0.5, 0.5]}]\' \\\n'
+ ' --outdir=/out\n\n'
'For more information about input file format, please see:\n'
'https://www.tensorflow.org/programmers_guide/saved_model_cli\n')
parser_run = subparsers.add_parser(
exp_out = """MetaGraphDef with tag-set: 'serve' contains the following SignatureDefs:
signature_def['classify_x2_to_y3']:
-The given SavedModel SignatureDef contains the following input(s):
-inputs['inputs'] tensor_info:
- dtype: DT_FLOAT
- shape: (-1, 1)
- name: x2:0
-The given SavedModel SignatureDef contains the following output(s):
-outputs['scores'] tensor_info:
- dtype: DT_FLOAT
- shape: (-1, 1)
- name: y3:0
-Method name is: tensorflow/serving/classify
+ The given SavedModel SignatureDef contains the following input(s):
+ inputs['inputs'] tensor_info:
+ dtype: DT_FLOAT
+ shape: (-1, 1)
+ name: x2:0
+ The given SavedModel SignatureDef contains the following output(s):
+ outputs['scores'] tensor_info:
+ dtype: DT_FLOAT
+ shape: (-1, 1)
+ name: y3:0
+ Method name is: tensorflow/serving/classify
signature_def['classify_x_to_y']:
-The given SavedModel SignatureDef contains the following input(s):
-inputs['inputs'] tensor_info:
- dtype: DT_STRING
- shape: unknown_rank
- name: tf_example:0
-The given SavedModel SignatureDef contains the following output(s):
-outputs['scores'] tensor_info:
- dtype: DT_FLOAT
- shape: (-1, 1)
- name: y:0
-Method name is: tensorflow/serving/classify
+ The given SavedModel SignatureDef contains the following input(s):
+ inputs['inputs'] tensor_info:
+ dtype: DT_STRING
+ shape: unknown_rank
+ name: tf_example:0
+ The given SavedModel SignatureDef contains the following output(s):
+ outputs['scores'] tensor_info:
+ dtype: DT_FLOAT
+ shape: (-1, 1)
+ name: y:0
+ Method name is: tensorflow/serving/classify
signature_def['regress_x2_to_y3']:
-The given SavedModel SignatureDef contains the following input(s):
-inputs['inputs'] tensor_info:
- dtype: DT_FLOAT
- shape: (-1, 1)
- name: x2:0
-The given SavedModel SignatureDef contains the following output(s):
-outputs['outputs'] tensor_info:
- dtype: DT_FLOAT
- shape: (-1, 1)
- name: y3:0
-Method name is: tensorflow/serving/regress
+ The given SavedModel SignatureDef contains the following input(s):
+ inputs['inputs'] tensor_info:
+ dtype: DT_FLOAT
+ shape: (-1, 1)
+ name: x2:0
+ The given SavedModel SignatureDef contains the following output(s):
+ outputs['outputs'] tensor_info:
+ dtype: DT_FLOAT
+ shape: (-1, 1)
+ name: y3:0
+ Method name is: tensorflow/serving/regress
signature_def['regress_x_to_y']:
-The given SavedModel SignatureDef contains the following input(s):
-inputs['inputs'] tensor_info:
- dtype: DT_STRING
- shape: unknown_rank
- name: tf_example:0
-The given SavedModel SignatureDef contains the following output(s):
-outputs['outputs'] tensor_info:
- dtype: DT_FLOAT
- shape: (-1, 1)
- name: y:0
-Method name is: tensorflow/serving/regress
+ The given SavedModel SignatureDef contains the following input(s):
+ inputs['inputs'] tensor_info:
+ dtype: DT_STRING
+ shape: unknown_rank
+ name: tf_example:0
+ The given SavedModel SignatureDef contains the following output(s):
+ outputs['outputs'] tensor_info:
+ dtype: DT_FLOAT
+ shape: (-1, 1)
+ name: y:0
+ Method name is: tensorflow/serving/regress
signature_def['regress_x_to_y2']:
-The given SavedModel SignatureDef contains the following input(s):
-inputs['inputs'] tensor_info:
- dtype: DT_STRING
- shape: unknown_rank
- name: tf_example:0
-The given SavedModel SignatureDef contains the following output(s):
-outputs['outputs'] tensor_info:
- dtype: DT_FLOAT
- shape: (-1, 1)
- name: y2:0
-Method name is: tensorflow/serving/regress
+ The given SavedModel SignatureDef contains the following input(s):
+ inputs['inputs'] tensor_info:
+ dtype: DT_STRING
+ shape: unknown_rank
+ name: tf_example:0
+ The given SavedModel SignatureDef contains the following output(s):
+ outputs['outputs'] tensor_info:
+ dtype: DT_FLOAT
+ shape: (-1, 1)
+ name: y2:0
+ Method name is: tensorflow/serving/regress
signature_def['serving_default']:
-The given SavedModel SignatureDef contains the following input(s):
-inputs['x'] tensor_info:
- dtype: DT_FLOAT
- shape: (-1, 1)
- name: x:0
-The given SavedModel SignatureDef contains the following output(s):
-outputs['y'] tensor_info:
- dtype: DT_FLOAT
- shape: (-1, 1)
- name: y:0
-Method name is: tensorflow/serving/predict"""
+ The given SavedModel SignatureDef contains the following input(s):
+ inputs['x'] tensor_info:
+ dtype: DT_FLOAT
+ shape: (-1, 1)
+ name: x:0
+ The given SavedModel SignatureDef contains the following output(s):
+ outputs['y'] tensor_info:
+ dtype: DT_FLOAT
+ shape: (-1, 1)
+ name: y:0
+ Method name is: tensorflow/serving/predict"""
# pylint: enable=line-too-long
+ self.maxDiff = None # Produce a useful error msg if the comparison fails
self.assertMultiLineEqual(output, exp_out)
self.assertEqual(err.getvalue().strip(), '')
output = out.getvalue().strip()
expected_output = (
'The given SavedModel SignatureDef contains the following input(s):\n'
- 'inputs[\'x\'] tensor_info:\n'
- ' dtype: DT_FLOAT\n shape: (-1, 1)\n name: x:0\n'
+ ' inputs[\'x\'] tensor_info:\n'
+ ' dtype: DT_FLOAT\n shape: (-1, 1)\n name: x:0\n'
'The given SavedModel SignatureDef contains the following output(s):\n'
- 'outputs[\'y\'] tensor_info:\n'
- ' dtype: DT_FLOAT\n shape: (-1, 1)\n name: y:0\n'
+ ' outputs[\'y\'] tensor_info:\n'
+ ' dtype: DT_FLOAT\n shape: (-1, 1)\n name: y:0\n'
'Method name is: tensorflow/serving/predict')
self.assertEqual(output, expected_output)
self.assertEqual(err.getvalue().strip(), '')
name: Name of the operation.
"""
base_type = variable.dtype.base_dtype
- with ops.colocate_with(variable.op):
+ # Do not colocate with variable since RestoreV2 op only runs on CPU and
+ # colocation will force variable (and other ops that colocate with variable)
+ # to be on CPU as well. It is okay to place the variable's initializer op on
+ # CPU since it will only be run once at the start.
+ with ops.device(variable.device), ops.device("/cpu:0"):
restore_op = io_ops.restore_v2(
ckpt_file, [tensor_name], [slice_spec], [base_type], name=name)[0]
if isinstance(variable, resource_variable_ops.ResourceVariable):
checkpoint_utils.init_from_checkpoint(checkpoint_dir,
{"useful_scope/": "useful_scope/"})
- self.assertEqual(my4._initializer_op.op.inputs[1].device, "/job:ps")
+ # initializer runs on the same task but always on CPU.
+ self.assertEqual(my4._initializer_op.op.inputs[1].device,
+ "/job:ps/device:CPU:0")
def testInitFromRootCheckpoint(self):
checkpoint_dir = self.get_temp_dir()
argspec: "args=[\'data\', \'segment_ids\', \'num_segments\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
+ name: "unsorted_segment_min"
+ argspec: "args=[\'data\', \'segment_ids\', \'num_segments\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "unsorted_segment_prod"
+ argspec: "args=[\'data\', \'segment_ids\', \'num_segments\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
+ name: "unsorted_segment_sqrt_n"
+ argspec: "args=[\'data\', \'segment_ids\', \'num_segments\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
+ }
+ member_method {
name: "unsorted_segment_sum"
argspec: "args=[\'data\', \'segment_ids\', \'num_segments\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
rm /this_is_writable_file_system
fi
+if [ -n "${CI_BUILD_USER_FORCE_BADNAME}" ]; then
+ ADDUSER_OPTS="--force-badname"
+fi
+
getent group "${CI_BUILD_GID}" || addgroup --gid "${CI_BUILD_GID}" "${CI_BUILD_GROUP}"
-getent passwd "${CI_BUILD_UID}" || adduser --gid "${CI_BUILD_GID}" --uid "${CI_BUILD_UID}" \
+getent passwd "${CI_BUILD_UID}" || adduser ${ADDUSER_OPTS} \
+ --gid "${CI_BUILD_GID}" --uid "${CI_BUILD_UID}" \
--gecos "${CI_BUILD_USER} (generated by with_the_same_user script)" \
--disabled-password --home "${CI_BUILD_HOME}" --quiet "${CI_BUILD_USER}"
usermod -a -G sudo "${CI_BUILD_USER}"
# ==============================================================================
# Select bazel version.
-BAZEL_VERSION="0.10.0"
+BAZEL_VERSION="0.11.0"
set +e
local_bazel_ver=$(bazel version 2>&1 | grep -i label | awk '{print $3}')
RUN echo "build --spawn_strategy=standalone --genrule_strategy=standalone" \
>>/etc/bazel.bazelrc
# Install the most recent bazel release.
-ENV BAZEL_VERSION 0.8.0
+ENV BAZEL_VERSION 0.11.0
WORKDIR /
RUN mkdir /bazel && \
cd /bazel && \
RUN echo "build --spawn_strategy=standalone --genrule_strategy=standalone" \
>>/etc/bazel.bazelrc
# Install the most recent bazel release.
-ENV BAZEL_VERSION 0.8.0
+ENV BAZEL_VERSION 0.11.0
WORKDIR /
RUN mkdir /bazel && \
cd /bazel && \
"//tensorflow/core:tensorflow",
"//tensorflow/contrib/rnn:gru_ops_op_lib",
"//tensorflow/contrib/rnn:lstm_ops_op_lib",
+ "//tensorflow/core/kernels:quantization_utils",
] + if_not_windows([
- "//tensorflow/core/kernels:quantized_ops",
"//tensorflow/core/kernels:remote_fused_graph_rewriter_transform",
"//tensorflow/core/kernels/hexagon:hexagon_rewriter_transform",
]),
// inputs which are referenced with "^tensor_name".
// See node_def.proto for more details.
Status RemoveControlDependencies(const GraphDef& input_graph_def,
- const TransformFuncContext& context,
- GraphDef* output_graph_def) {
- output_graph_def->Clear();
- for (const NodeDef& node : input_graph_def.node()) {
- NodeDef* new_node = output_graph_def->mutable_node()->Add();
- *new_node = node;
- new_node->clear_input();
- for (const auto& input : node.input()) {
- if (input[0] != '^') {
- new_node->add_input(input);
- }
+ const TransformFuncContext& context,
+ GraphDef* output_graph_def) {
+ output_graph_def->Clear();
+ for (const NodeDef& node : input_graph_def.node()) {
+ NodeDef* new_node = output_graph_def->mutable_node()->Add();
+ *new_node = node;
+ new_node->clear_input();
+ for (const auto& input : node.input()) {
+ if (input[0] != '^') {
+ new_node->add_input(input);
+ }
+ }
}
- }
- return Status::OK();
+ return Status::OK();
}
-REGISTER_GRAPH_TRANSFORM("remove_control_dependencies",
- RemoveControlDependencies);
+REGISTER_GRAPH_TRANSFORM("remove_control_dependencies", RemoveControlDependencies);
} // namespace graph_transforms
} // namespace tensorflow
":cheaders",
":clib",
":clicenses",
+ ":eager_cheaders",
],
)
name = "cheaders",
files = [
"//tensorflow/c:headers",
- "//tensorflow/c/eager:headers",
],
package_dir = "include/tensorflow/c",
# Mark as "manual" till
)
pkg_tar(
+ name = "eager_cheaders",
+ files = [
+ "//tensorflow/c/eager:headers",
+ ],
+ package_dir = "include/tensorflow/c/eager",
+ # Mark as "manual" till
+ # https://github.com/bazelbuild/bazel/issues/2352
+ # and https://github.com/bazelbuild/bazel/issues/1580
+ # are resolved, otherwise these rules break when built
+ # with Python 3.
+ tags = ["manual"],
+)
+
+pkg_tar(
name = "clib",
files = ["//tensorflow:libtensorflow.so"],
package_dir = "lib",
projects / platform / upstream / tensorflow.git / commitdiff